From 5457f482905419ff966a237018cf2ce03fb21506 Mon Sep 17 00:00:00 2001 From: Shivi Jain Date: Mon, 14 Sep 2026 20:50:28 +0530 Subject: [PATCH 01/13] fix(auth): load team membership once per request and skip prisma on an L1 hit common_checks was querying get_team_membership twice, and DualCache awaited Redis SET on the auth path, so LRU eviction plus a hung Redis write showed up as two postgres spans --- litellm/proxy/auth/auth_checks.py | 313 ++++++++++++++---- .../proxy/auth/test_auth_checks.py | 127 +++++++ 2 files changed, 369 insertions(+), 71 deletions(-) diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 576585ee9a3..abcade6b1c8 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -327,6 +327,9 @@ _safe_json_loads_obj: Final = _typed_json_loads(safe_json_loads) last_db_access_time: Final = LimitedSizeOrderedDict(max_size=100) db_cache_expiry: Final = DEFAULT_IN_MEMORY_TTL # refresh every 5s +_TEAM_MEMBERSHIP_CACHE_MISS: Final = object() +_team_membership_inflight: dict[str, asyncio.Task[LiteLLM_TeamMembership | None]] = {} + all_routes: Final = LiteLLMRoutes.openai_routes.value + LiteLLMRoutes.management_routes.value @@ -873,6 +876,21 @@ async def common_checks( """ from litellm.proxy.proxy_server import prisma_client, user_api_key_cache + # One membership read for model-access, access-group attribution, and + # member-budget. Each used to call get_team_membership independently; + # DualCache Redis SET/GET on every miss made those look like two Postgres spans. + loaded_team_membership: LiteLLM_TeamMembership | None = None + team_membership_loaded = False + if team_object is not None and valid_token is not None and valid_token.user_id is not None: + loaded_team_membership = await get_team_membership( + user_id=valid_token.user_id, + team_id=team_object.team_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + team_membership_loaded = True + _model: Final[str | list[str] | None] = get_model_from_request( request_data=request_body, route=route, @@ -936,6 +954,8 @@ async def common_checks( prisma_client=prisma_client, user_api_key_cache=user_api_key_cache, proxy_logging_obj=proxy_logging_obj, + team_membership=loaded_team_membership, + team_membership_loaded=team_membership_loaded, ) # Require trace id for agent keys when agent has require_trace_id_on_calls_by_agent @@ -987,6 +1007,8 @@ async def common_checks( prisma_client=prisma_client, user_api_key_cache=user_api_key_cache, proxy_logging_obj=proxy_logging_obj, + team_membership=loaded_team_membership, + team_membership_loaded=team_membership_loaded, ) # Run before apply_key_tags_pre_auth injects key metadata.tags into request_body. @@ -1096,6 +1118,8 @@ async def common_checks( prisma_client=prisma_client, user_api_key_cache=user_api_key_cache, proxy_logging_obj=proxy_logging_obj, + team_membership=loaded_team_membership, + team_membership_loaded=team_membership_loaded, ), _check_end_user_budget(end_user_obj=end_user_object, route=route) if end_user_object is not None and end_user_object.litellm_budget_table is not None @@ -2141,7 +2165,142 @@ async def get_tag_object( return tag_objects.get(tag_name) +def _debug_team_membership_log(hypothesis_id: str, message: str, data: dict[str, object]) -> None: + # #region agent log + try: + import json as _json + + with open("/Users/shijain/genai-apps/genai-proxy/.cursor/debug-86534f.log", "a", encoding="utf-8") as _f: + _f.write( + _json.dumps( + { + "sessionId": "86534f", + "timestamp": int(time.time() * 1000), + "location": "auth_checks.py:get_team_membership", + "message": message, + "hypothesisId": hypothesis_id, + "data": data, + } + ) + + "\n" + ) + except Exception: + pass + # #endregion + + +def _membership_from_cached_payload(cached: object) -> LiteLLM_TeamMembership | None | object: + """Decode a DualCache payload. ``_TEAM_MEMBERSHIP_CACHE_MISS`` means try the next tier.""" + if cached is None: + return _TEAM_MEMBERSHIP_CACHE_MISS + if cached == NO_TEAM_MEMBERSHIP_SENTINEL: + return None + cached_membership: Final = CacheCodec.deserialize(cached, model_type=LiteLLM_TeamMembership) + return cached_membership if cached_membership is not None else _TEAM_MEMBERSHIP_CACHE_MISS + + +async def _populate_team_membership_cache( + user_api_key_cache: UserApiKeyCache, + key: str, + value: object, + *, + model_type: type[LiteLLM_TeamMembership] | None = None, + ttl: float | None = None, +) -> None: + """Await in-memory write; replicate to Redis off the auth await path.""" + kwargs: dict[str, object] = {} + if model_type is not None: + kwargs["model_type"] = model_type + if ttl is not None: + kwargs["ttl"] = ttl + await user_api_key_cache.async_set_cache(key=key, value=value, local_only=True, **kwargs) + + async def _replicate_to_redis() -> None: + try: + await user_api_key_cache.async_set_cache(key=key, value=value, **kwargs) + except Exception: + return + + asyncio.create_task(_replicate_to_redis()) + + @log_db_metrics +async def _fetch_team_membership_from_db( + user_id: str, + team_id: str, + prisma_client: PrismaClient, + user_api_key_cache: UserApiKeyCache, + parent_otel_span: Span | None = None, + proxy_logging_obj: ProxyLogging | None = None, +) -> LiteLLM_TeamMembership | None: + """Prisma read + L1 populate. Decorated so cache hits on ``get_team_membership`` are not postgres spans.""" + _ = parent_otel_span, proxy_logging_obj + response: Final = await _dictable_table(TeamMembershipRepository(prisma_client)).find_unique( + where={"user_id_team_id": {"user_id": user_id, "team_id": team_id}}, + include={"litellm_budget_table": True}, + ) + _key: Final = team_membership_reservation_cache_key(user_id=user_id, team_id=team_id) + if response is None: + await _populate_team_membership_cache( + user_api_key_cache, + _key, + NO_TEAM_MEMBERSHIP_SENTINEL, + ttl=get_management_object_ttl(user_api_key_cache), + ) + return None + + membership: Final = LiteLLM_TeamMembership.model_validate(response.dict()) + await _populate_team_membership_cache( + user_api_key_cache, + _key, + membership, + model_type=LiteLLM_TeamMembership, + ) + return membership + + +async def _load_team_membership_on_cache_miss( + user_id: str, + team_id: str, + cache_key: str, + prisma_client: PrismaClient, + user_api_key_cache: UserApiKeyCache, + parent_otel_span: Span | None, + proxy_logging_obj: ProxyLogging | None, +) -> LiteLLM_TeamMembership | None: + try: + redis_cached: Final[object] = await user_api_key_cache.async_get_cache(key=cache_key) + redis_membership: Final = _membership_from_cached_payload(redis_cached) + if redis_membership is not _TEAM_MEMBERSHIP_CACHE_MISS: + # #region agent log + _debug_team_membership_log( + "H3", + "membership redis hit after l1 miss", + {"prisma": False, "source": "redis"}, + ) + # #endregion + return cast(LiteLLM_TeamMembership | None, redis_membership) + + # #region agent log + _debug_team_membership_log("H1", "membership prisma fetch", {"prisma": True}) + # #endregion + return await _fetch_team_membership_from_db( + user_id=user_id, + team_id=team_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=parent_otel_span, + proxy_logging_obj=proxy_logging_obj, + ) + except Exception: + verbose_proxy_logger.exception( + "Error getting team membership for user_id: %s, team_id: %s", + user_id, + team_id, + ) + return None + + async def get_team_membership( user_id: str, team_id: str, @@ -2155,54 +2314,43 @@ async def get_team_membership( Do a isolated check for team membership vs. doing a combined key + team + user + team-membership check, as key might come in frequently for different users/teams. Larger call will slowdown query time. This way we get to cache the constant (key/team/user info) and only update based on the changing value (team membership). """ - from litellm.proxy._types import LiteLLM_TeamMembership - - if prisma_client is None: - raise Exception("No db connected") - if user_id is None or team_id is None: return None _key: Final = team_membership_reservation_cache_key(user_id=user_id, team_id=team_id) - # check if in cache - cached: Final[object] = await user_api_key_cache.async_get_cache(key=_key) - if cached == NO_TEAM_MEMBERSHIP_SENTINEL: - return None - cached_membership_obj: Final = CacheCodec.deserialize(cached, model_type=LiteLLM_TeamMembership) - if cached_membership_obj is not None: - return cached_membership_obj + l1_cached: Final[object] = await user_api_key_cache.async_get_cache(key=_key, local_only=True) + l1_membership: Final = _membership_from_cached_payload(l1_cached) + if l1_membership is not _TEAM_MEMBERSHIP_CACHE_MISS: + # #region agent log + _debug_team_membership_log("H4", "membership l1 hit", {"prisma": False, "source": "l1"}) + # #endregion + return cast(LiteLLM_TeamMembership | None, l1_membership) - # else, check db - try: - response: Final = await _dictable_table(TeamMembershipRepository(prisma_client)).find_unique( - where={"user_id_team_id": {"user_id": user_id, "team_id": team_id}}, - include={"litellm_budget_table": True}, + inflight: Final = _team_membership_inflight.get(_key) + if inflight is not None: + # #region agent log + _debug_team_membership_log("H5", "membership coalesced waiter", {"prisma": False, "coalesced": True}) + # #endregion + return await inflight + + if prisma_client is None: + raise Exception("No db connected") + + task: Final = asyncio.ensure_future( + _load_team_membership_on_cache_miss( + user_id=user_id, + team_id=team_id, + cache_key=_key, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=parent_otel_span, + proxy_logging_obj=proxy_logging_obj, ) - - if response is None: - await user_api_key_cache.async_set_cache( - key=_key, - value=NO_TEAM_MEMBERSHIP_SENTINEL, - ttl=get_management_object_ttl(user_api_key_cache), - ) - return None - - _response: Final = LiteLLM_TeamMembership.model_validate(response.dict()) - await user_api_key_cache.async_set_cache( - key=_key, - value=_response, - model_type=LiteLLM_TeamMembership, - ) - - return _response - except Exception: - verbose_proxy_logger.exception( - "Error getting team membership for user_id: %s, team_id: %s", - user_id, - team_id, - ) - return None + ) + _team_membership_inflight[_key] = task + task.add_done_callback(lambda _t, k=_key: _team_membership_inflight.pop(k, None)) + return await task def model_in_access_group(model: str, team_models: list[str] | None, llm_router: Router | None) -> bool: @@ -4122,18 +4270,21 @@ async def _team_member_granted_models( prisma_client: PrismaClient, user_api_key_cache: UserApiKeyCache, proxy_logging_obj: ProxyLogging, + team_membership: LiteLLM_TeamMembership | None = None, + team_membership_loaded: bool = False, ) -> Sequence[str]: """The member's own ``allowed_models`` scope; empty when the member is not narrowed below the team.""" if team_object is None or valid_token.user_id is None: return () - team_membership: Final = await get_team_membership( - user_id=valid_token.user_id, - team_id=team_object.team_id, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=proxy_logging_obj, - ) + if not team_membership_loaded: + team_membership = await get_team_membership( + user_id=valid_token.user_id, + team_id=team_object.team_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) return () if team_membership is None else _member_allowed_models(team_membership) @@ -4169,6 +4320,8 @@ async def _granted_model_lists( prisma_client: PrismaClient, user_api_key_cache: UserApiKeyCache, proxy_logging_obj: ProxyLogging, + team_membership: LiteLLM_TeamMembership | None = None, + team_membership_loaded: bool = False, ) -> tuple[Sequence[str], ...]: """One model allowlist per level that participates in authorizing the request.""" return ( @@ -4180,6 +4333,8 @@ async def _granted_model_lists( prisma_client=prisma_client, user_api_key_cache=user_api_key_cache, proxy_logging_obj=proxy_logging_obj, + team_membership=team_membership, + team_membership_loaded=team_membership_loaded, ), project_object.models if project_object is not None else (), await _org_granted_models( @@ -4274,6 +4429,8 @@ async def collect_matched_model_access_groups( prisma_client: PrismaClient | None, user_api_key_cache: UserApiKeyCache, proxy_logging_obj: ProxyLogging, + team_membership: LiteLLM_TeamMembership | None = None, + team_membership_loaded: bool = False, ) -> tuple[str, ...]: """ The budgeted model access groups that authorized this request, sorted and deduplicated. @@ -4319,6 +4476,8 @@ async def collect_matched_model_access_groups( prisma_client=prisma_client, user_api_key_cache=user_api_key_cache, proxy_logging_obj=proxy_logging_obj, + team_membership=team_membership, + team_membership_loaded=team_membership_loaded, ) for granted_model in granted_models ) @@ -4334,6 +4493,8 @@ async def stamp_matched_model_access_groups( prisma_client: PrismaClient | None, user_api_key_cache: UserApiKeyCache, proxy_logging_obj: ProxyLogging, + team_membership: LiteLLM_TeamMembership | None = None, + team_membership_loaded: bool = False, ) -> tuple[str, ...]: """Record the groups that authorized this request on its auth object, for the post-call spend writer and the reservation counters, and hand them back for the budget check.""" @@ -4350,6 +4511,8 @@ async def stamp_matched_model_access_groups( prisma_client=prisma_client, user_api_key_cache=user_api_key_cache, proxy_logging_obj=proxy_logging_obj, + team_membership=team_membership, + team_membership_loaded=team_membership_loaded, ) except Exception as e: # noqa: BLE001 # fail-safe: attribution is spend telemetry, it must never break auth verbose_proxy_logger.debug("model access group attribution failed: %s", e) @@ -5152,6 +5315,8 @@ async def _check_team_member_budget( prisma_client: PrismaClient | None, user_api_key_cache: UserApiKeyCache, proxy_logging_obj: ProxyLogging, + team_membership: LiteLLM_TeamMembership | None = None, + team_membership_loaded: bool = False, ): """Check if team member is over their max budget within the team.""" if ( @@ -5160,23 +5325,25 @@ async def _check_team_member_budget( and valid_token is not None and valid_token.user_id is not None ): - team_membership: Final = await get_team_membership( - user_id=valid_token.user_id, - team_id=team_object.team_id, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=proxy_logging_obj, - ) + if not team_membership_loaded: + team_membership = await get_team_membership( + user_id=valid_token.user_id, + team_id=team_object.team_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + loaded_membership = team_membership # Per-member override wins; otherwise fall back to the team-level # default configured via team.metadata["team_member_budget_id"]. team_member_budget: float | None = None if ( - team_membership is not None - and team_membership.litellm_budget_table is not None - and team_membership.litellm_budget_table.max_budget is not None + loaded_membership is not None + and loaded_membership.litellm_budget_table is not None + and loaded_membership.litellm_budget_table.max_budget is not None ): - team_member_budget = team_membership.litellm_budget_table.max_budget + team_member_budget = loaded_membership.litellm_budget_table.max_budget else: default_budget_id: Final = (team_object.metadata or {}).get("team_member_budget_id") if isinstance(default_budget_id, str): @@ -5195,7 +5362,7 @@ async def _check_team_member_budget( team_member_budget = default_budget.max_budget if team_member_budget is not None: - team_member_spend = (team_membership.spend if team_membership is not None else 0.0) or 0.0 + team_member_spend = (loaded_membership.spend if loaded_membership is not None else 0.0) or 0.0 # Read from cross-pod counter (Redis-first) if available from litellm.proxy.proxy_server import get_current_spend @@ -5224,6 +5391,8 @@ async def _check_team_member_model_access( prisma_client: Optional["PrismaClient"], user_api_key_cache: UserApiKeyCache, proxy_logging_obj: ProxyLogging, + team_membership: LiteLLM_TeamMembership | None = None, + team_membership_loaded: bool = False, ) -> None: """ Check if a team member's per-member model scope allows access to the requested model. @@ -5234,22 +5403,24 @@ async def _check_team_member_model_access( if valid_token.user_id is None or team_object.team_id is None: return - team_membership: Final = await get_team_membership( - user_id=valid_token.user_id, - team_id=team_object.team_id, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=proxy_logging_obj, - ) + if not team_membership_loaded: + team_membership = await get_team_membership( + user_id=valid_token.user_id, + team_id=team_object.team_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + loaded_membership = team_membership if ( - team_membership is None - or team_membership.litellm_budget_table is None - or not team_membership.litellm_budget_table.allowed_models + loaded_membership is None + or loaded_membership.litellm_budget_table is None + or not loaded_membership.litellm_budget_table.allowed_models ): return # no per-member restriction — inherit team-level check - member_allowed_models: Final[list[str]] = team_membership.litellm_budget_table.allowed_models + member_allowed_models: Final[list[str]] = loaded_membership.litellm_budget_table.allowed_models try: _can_object_call_model( model=model, diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index ea2c925212d..755fe9efe2c 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -6454,6 +6454,133 @@ async def test_get_team_membership_reads_sentinel_as_no_membership_not_a_model() mock_prisma_client.db.litellm_teammembership.find_unique.assert_not_awaited() +@pytest.mark.asyncio +async def test_get_team_membership_coalesces_parallel_db_fetches(): + """Concurrent misses for the same member must share one Prisma round-trip.""" + from litellm.proxy.auth.auth_checks import get_team_membership + + started = asyncio.Event() + release = asyncio.Event() + membership_row = MagicMock() + membership_row.dict = lambda: {"user_id": "u-parallel", "team_id": "t-parallel", "spend": 1.0} + + async def _slow_find_unique(*args, **kwargs): + started.set() + await release.wait() + return membership_row + + mock_prisma_client = MagicMock() + mock_prisma_client.db.litellm_teammembership.find_unique = AsyncMock(side_effect=_slow_find_unique) + cache = UserApiKeyCache() + + async def _load(): + return await get_team_membership( + user_id="u-parallel", + team_id="t-parallel", + prisma_client=mock_prisma_client, + user_api_key_cache=cache, + ) + + first = asyncio.create_task(_load()) + second = asyncio.create_task(_load()) + await started.wait() + await asyncio.sleep(0) + release.set() + results = await asyncio.gather(first, second) + + assert results[0] is not None and results[1] is not None + assert results[0].user_id == "u-parallel" + assert results[1].user_id == "u-parallel" + mock_prisma_client.db.litellm_teammembership.find_unique.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_get_team_membership_returns_before_redis_set_completes(): + """Auth must not wait on DualCache Redis SET; L1 is enough for the next lookup.""" + from litellm.proxy.auth.auth_checks import get_team_membership + + membership_row = MagicMock() + membership_row.dict = lambda: {"user_id": "u-redis", "team_id": "t-redis", "spend": 2.0} + + hang_redis_set = asyncio.Event() + + async def _hanging_redis_set(*args, **kwargs): + await hang_redis_set.wait() + + redis_cache = MagicMock() + redis_cache.async_get_cache = AsyncMock(return_value=None) + redis_cache.async_set_cache = AsyncMock(side_effect=_hanging_redis_set) + + cache = UserApiKeyCache(redis_cache=redis_cache) + mock_prisma_client = MagicMock() + mock_prisma_client.db.litellm_teammembership.find_unique = AsyncMock(return_value=membership_row) + + first = await asyncio.wait_for( + get_team_membership( + user_id="u-redis", + team_id="t-redis", + prisma_client=mock_prisma_client, + user_api_key_cache=cache, + ), + timeout=0.5, + ) + second = await get_team_membership( + user_id="u-redis", + team_id="t-redis", + prisma_client=mock_prisma_client, + user_api_key_cache=cache, + ) + + hang_redis_set.set() + await asyncio.sleep(0) + + assert first is not None and second is not None + assert first.spend == 2.0 + assert second.spend == 2.0 + mock_prisma_client.db.litellm_teammembership.find_unique.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_common_checks_calls_get_team_membership_once_per_request(): + """Model-access, attribution, and member-budget must reuse one membership load.""" + from fastapi import Request + + from litellm.proxy.auth.auth_checks import common_checks + + team = LiteLLM_TeamTable(team_id="t-once") + token = UserAPIKeyAuth(token="k-once", user_id="u-once", team_id="t-once", models=["gpt-4o-mini"]) + membership = MagicMock() + membership.litellm_budget_table = None + membership.spend = 0.0 + + with ( + patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), + patch("litellm.proxy.proxy_server.user_api_key_cache", UserApiKeyCache()), + patch( + "litellm.proxy.auth.auth_checks.get_team_membership", + new_callable=AsyncMock, + return_value=membership, + ) as load_membership, + patch("litellm.proxy.proxy_server.get_current_spend", new_callable=AsyncMock, return_value=0.0), + ): + result = await common_checks( + request_body={"model": "gpt-4o-mini", "messages": [{"role": "user", "content": "hi"}]}, + team_object=team, + user_object=LiteLLM_UserTable(user_id="u-once"), + end_user_object=None, + global_proxy_spend=None, + general_settings={}, + route="/chat/completions", + llm_router=None, + proxy_logging_obj=MagicMock(), + valid_token=token, + request=MagicMock(spec=Request), + ) + + assert result is True + assert load_membership.await_count == 1 + + @pytest.mark.asyncio async def test_invalidate_team_member_spend_state_evicts_the_negative_cache_sentinel(): """ From 70ddc7e4928b19fe2d73cc97ff707e0919a8bceb Mon Sep 17 00:00:00 2001 From: Shivi Jain Date: Mon, 14 Sep 2026 21:28:08 +0530 Subject: [PATCH 02/13] fix(auth): load team membership once per request and skip prisma on an L1 hit common_checks was querying get_team_membership twice, and DualCache awaited Redis SET on the auth path, so LRU eviction plus a hung Redis write showed up as two postgres spans --- litellm/constants.py | 7 + litellm/proxy/auth/auth_checks.py | 163 +++++++++++------- .../proxy/auth/test_auth_checks.py | 160 +++++++++++++++++ 3 files changed, 265 insertions(+), 65 deletions(-) diff --git a/litellm/constants.py b/litellm/constants.py index 5751e6e46af..5c7a02d0743 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -2039,3 +2039,10 @@ BATCH_ENQUEUED_TOKEN_LIMIT_METADATA_KEY: Final = "batch_enqueued_token_limit" # Shared read-only empty mapping, for defaulting optional Mapping parameters without # constructing a fresh mutable dict at each call site. EMPTY_MAPPING: Final = MappingProxyType({}) + + +class TeamMembershipCacheMiss: + __slots__ = () + + +TEAM_MEMBERSHIP_CACHE_MISS: Final = TeamMembershipCacheMiss() diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index abcade6b1c8..0baa79fb95a 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -35,6 +35,8 @@ from litellm.constants import ( MODEL_ACCESS_GROUP_REGISTRY_MAX_SIZE, REGISTRY_ERROR_NEGATIVE_CACHE_TTL, TAG_REGISTRY_MAX_SIZE, + TEAM_MEMBERSHIP_CACHE_MISS, + TeamMembershipCacheMiss, ) from litellm.litellm_core_utils.dd_tracing import tracer from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider @@ -327,12 +329,31 @@ _safe_json_loads_obj: Final = _typed_json_loads(safe_json_loads) last_db_access_time: Final = LimitedSizeOrderedDict(max_size=100) db_cache_expiry: Final = DEFAULT_IN_MEMORY_TTL # refresh every 5s -_TEAM_MEMBERSHIP_CACHE_MISS: Final = object() -_team_membership_inflight: dict[str, asyncio.Task[LiteLLM_TeamMembership | None]] = {} +_TEAM_MEMBERSHIP_INFLIGHT_MAX: Final = 10000 +_team_membership_inflight: Final = LimitedSizeOrderedDict(max_size=_TEAM_MEMBERSHIP_INFLIGHT_MAX) +_team_membership_write_epoch: Final = LimitedSizeOrderedDict(max_size=_TEAM_MEMBERSHIP_INFLIGHT_MAX) all_routes: Final = LiteLLMRoutes.openai_routes.value + LiteLLMRoutes.management_routes.value +def _membership_write_epoch(key: str) -> int: + cached: Final[object] = _team_membership_write_epoch.get(key, 0) + return cached if isinstance(cached, int) else 0 + + +def _bump_membership_write_epoch(key: str) -> None: + _team_membership_write_epoch[key] = _membership_write_epoch(key) + 1 + + +def _membership_from_shared_load(result: object) -> LiteLLM_TeamMembership | None: + if result is None or isinstance(result, LiteLLM_TeamMembership): + return result + raise HTTPException( + status_code=status.HTTP_503_SERVICE_UNAVAILABLE, + detail="Failed to load team membership", + ) + + def _log_budget_lookup_failure(entity: str, error: Exception) -> None: """ Log a warning when budget lookup fails; cache will not be populated. @@ -2165,38 +2186,39 @@ async def get_tag_object( return tag_objects.get(tag_name) -def _debug_team_membership_log(hypothesis_id: str, message: str, data: dict[str, object]) -> None: - # #region agent log - try: - import json as _json - - with open("/Users/shijain/genai-apps/genai-proxy/.cursor/debug-86534f.log", "a", encoding="utf-8") as _f: - _f.write( - _json.dumps( - { - "sessionId": "86534f", - "timestamp": int(time.time() * 1000), - "location": "auth_checks.py:get_team_membership", - "message": message, - "hypothesisId": hypothesis_id, - "data": data, - } - ) - + "\n" - ) - except Exception: - pass - # #endregion - - -def _membership_from_cached_payload(cached: object) -> LiteLLM_TeamMembership | None | object: - """Decode a DualCache payload. ``_TEAM_MEMBERSHIP_CACHE_MISS`` means try the next tier.""" +def _membership_from_cached_payload( + cached: object, +) -> LiteLLM_TeamMembership | None | TeamMembershipCacheMiss: if cached is None: - return _TEAM_MEMBERSHIP_CACHE_MISS + return TEAM_MEMBERSHIP_CACHE_MISS if cached == NO_TEAM_MEMBERSHIP_SENTINEL: return None cached_membership: Final = CacheCodec.deserialize(cached, model_type=LiteLLM_TeamMembership) - return cached_membership if cached_membership is not None else _TEAM_MEMBERSHIP_CACHE_MISS + return cached_membership if cached_membership is not None else TEAM_MEMBERSHIP_CACHE_MISS + + +async def _set_team_membership_cache_entry( + user_api_key_cache: UserApiKeyCache, + key: str, + value: object, + *, + local_only: bool, + model_type: type[LiteLLM_TeamMembership] | None, + ttl: float | None, +) -> None: + match (model_type is not None, ttl is not None): + case (False, False): + await user_api_key_cache.async_set_cache(key=key, value=value, local_only=local_only) + case (False, True): + await user_api_key_cache.async_set_cache(key=key, value=value, local_only=local_only, ttl=ttl) + case (True, False): + await user_api_key_cache.async_set_cache( + key=key, value=value, local_only=local_only, model_type=model_type + ) + case (True, True): + await user_api_key_cache.async_set_cache( + key=key, value=value, local_only=local_only, model_type=model_type, ttl=ttl + ) async def _populate_team_membership_cache( @@ -2207,17 +2229,33 @@ async def _populate_team_membership_cache( model_type: type[LiteLLM_TeamMembership] | None = None, ttl: float | None = None, ) -> None: - """Await in-memory write; replicate to Redis off the auth await path.""" - kwargs: dict[str, object] = {} - if model_type is not None: - kwargs["model_type"] = model_type - if ttl is not None: - kwargs["ttl"] = ttl - await user_api_key_cache.async_set_cache(key=key, value=value, local_only=True, **kwargs) + write_epoch: Final = _membership_write_epoch(key) + await _set_team_membership_cache_entry( + user_api_key_cache, + key, + value, + local_only=True, + model_type=model_type, + ttl=ttl, + ) + if _membership_write_epoch(key) != write_epoch: + await user_api_key_cache.async_delete_cache(key) + return async def _replicate_to_redis() -> None: try: - await user_api_key_cache.async_set_cache(key=key, value=value, **kwargs) + if _membership_write_epoch(key) != write_epoch: + return + await _set_team_membership_cache_entry( + user_api_key_cache, + key, + value, + local_only=False, + model_type=model_type, + ttl=ttl, + ) + if _membership_write_epoch(key) != write_epoch: + await user_api_key_cache.async_delete_cache(key) except Exception: return @@ -2271,19 +2309,9 @@ async def _load_team_membership_on_cache_miss( try: redis_cached: Final[object] = await user_api_key_cache.async_get_cache(key=cache_key) redis_membership: Final = _membership_from_cached_payload(redis_cached) - if redis_membership is not _TEAM_MEMBERSHIP_CACHE_MISS: - # #region agent log - _debug_team_membership_log( - "H3", - "membership redis hit after l1 miss", - {"prisma": False, "source": "redis"}, - ) - # #endregion - return cast(LiteLLM_TeamMembership | None, redis_membership) + if not isinstance(redis_membership, TeamMembershipCacheMiss): + return redis_membership - # #region agent log - _debug_team_membership_log("H1", "membership prisma fetch", {"prisma": True}) - # #endregion return await _fetch_team_membership_from_db( user_id=user_id, team_id=team_id, @@ -2292,13 +2320,18 @@ async def _load_team_membership_on_cache_miss( parent_otel_span=parent_otel_span, proxy_logging_obj=proxy_logging_obj, ) - except Exception: + except HTTPException: + raise + except Exception as e: verbose_proxy_logger.exception( "Error getting team membership for user_id: %s, team_id: %s", user_id, team_id, ) - return None + raise HTTPException( + status_code=status.HTTP_503_SERVICE_UNAVAILABLE, + detail="Failed to load team membership", + ) from e async def get_team_membership( @@ -2321,18 +2354,12 @@ async def get_team_membership( l1_cached: Final[object] = await user_api_key_cache.async_get_cache(key=_key, local_only=True) l1_membership: Final = _membership_from_cached_payload(l1_cached) - if l1_membership is not _TEAM_MEMBERSHIP_CACHE_MISS: - # #region agent log - _debug_team_membership_log("H4", "membership l1 hit", {"prisma": False, "source": "l1"}) - # #endregion - return cast(LiteLLM_TeamMembership | None, l1_membership) + if not isinstance(l1_membership, TeamMembershipCacheMiss): + return l1_membership - inflight: Final = _team_membership_inflight.get(_key) - if inflight is not None: - # #region agent log - _debug_team_membership_log("H5", "membership coalesced waiter", {"prisma": False, "coalesced": True}) - # #endregion - return await inflight + inflight: Final[object] = _team_membership_inflight.get(_key) + if isinstance(inflight, asyncio.Task): + return _membership_from_shared_load(await asyncio.shield(inflight)) if prisma_client is None: raise Exception("No db connected") @@ -2349,8 +2376,13 @@ async def get_team_membership( ) ) _team_membership_inflight[_key] = task - task.add_done_callback(lambda _t, k=_key: _team_membership_inflight.pop(k, None)) - return await task + + def _clear_inflight(_done: object) -> None: + if _team_membership_inflight.get(_key) is task: + _team_membership_inflight.pop(_key, None) + + task.add_done_callback(_clear_inflight) + return _membership_from_shared_load(await asyncio.shield(task)) def model_in_access_group(model: str, team_models: list[str] | None, llm_router: Router | None) -> bool: @@ -2861,6 +2893,7 @@ async def invalidate_team_member_spend_state( ttl=SPEND_DB_FLOOR_CACHE_TTL_SECONDS, ) + _bump_membership_write_epoch(team_membership_reservation_cache_key(user_id=user_id, team_id=team_id)) await evict_and_broadcast( cache_keys=( team_membership_auth_cache_key(team_id=team_id, user_id=user_id), diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index 755fe9efe2c..d16c75a4a3b 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -6581,6 +6581,166 @@ async def test_common_checks_calls_get_team_membership_once_per_request(): assert load_membership.await_count == 1 +@pytest.mark.asyncio +async def test_get_team_membership_db_error_raises_503_not_none(): + """A Prisma failure must fail closed as 503, not look like a missing membership row.""" + from fastapi import HTTPException + + from litellm.proxy.auth.auth_checks import get_team_membership + from litellm.proxy.common_utils.user_api_key_cache import team_membership_reservation_cache_key + + mock_prisma_client = MagicMock() + mock_prisma_client.db.litellm_teammembership.find_unique = AsyncMock(side_effect=RuntimeError("db down")) + cache = UserApiKeyCache() + + with pytest.raises(HTTPException) as exc: + await get_team_membership( + user_id="u-fail", + team_id="t-fail", + prisma_client=mock_prisma_client, + user_api_key_cache=cache, + ) + + cached = await cache.async_get_cache( + key=team_membership_reservation_cache_key(user_id="u-fail", team_id="t-fail") + ) + assert exc.value.status_code == status.HTTP_503_SERVICE_UNAVAILABLE + assert cached is None + + +@pytest.mark.asyncio +async def test_common_checks_does_not_skip_member_limits_when_membership_lookup_fails(): + """common_checks must not mark membership loaded-absent after a lookup error.""" + from fastapi import HTTPException, Request + + from litellm.proxy.auth.auth_checks import common_checks + + team = LiteLLM_TeamTable(team_id="t-fail-closed") + token = UserAPIKeyAuth( + token="k-fail-closed", + user_id="u-fail-closed", + team_id="t-fail-closed", + models=["gpt-4o-mini"], + ) + lookup_error = HTTPException( + status_code=status.HTTP_503_SERVICE_UNAVAILABLE, + detail="Failed to load team membership", + ) + + with ( + patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), + patch("litellm.proxy.proxy_server.user_api_key_cache", UserApiKeyCache()), + patch( + "litellm.proxy.auth.auth_checks.get_team_membership", + new_callable=AsyncMock, + side_effect=lookup_error, + ), + patch("litellm.proxy.proxy_server.get_current_spend", new_callable=AsyncMock, return_value=0.0), + ): + with pytest.raises(HTTPException) as exc: + await common_checks( + request_body={"model": "gpt-4o-mini", "messages": [{"role": "user", "content": "hi"}]}, + team_object=team, + user_object=LiteLLM_UserTable(user_id="u-fail-closed"), + end_user_object=None, + global_proxy_spend=None, + general_settings={}, + route="/chat/completions", + llm_router=None, + proxy_logging_obj=MagicMock(), + valid_token=token, + request=MagicMock(spec=Request), + ) + + assert exc.value.status_code == status.HTTP_503_SERVICE_UNAVAILABLE + + +@pytest.mark.asyncio +async def test_get_team_membership_waiter_cancel_does_not_cancel_shared_load(): + """Cancelling one coalesced waiter must not cancel the shared Prisma load.""" + from litellm.proxy.auth.auth_checks import get_team_membership + + started = asyncio.Event() + release = asyncio.Event() + membership_row = MagicMock() + membership_row.dict = lambda: {"user_id": "u-shield", "team_id": "t-shield", "spend": 1.0} + + async def _slow_find_unique(*args, **kwargs): + started.set() + await release.wait() + return membership_row + + mock_prisma_client = MagicMock() + mock_prisma_client.db.litellm_teammembership.find_unique = AsyncMock(side_effect=_slow_find_unique) + cache = UserApiKeyCache() + + async def _load(): + return await get_team_membership( + user_id="u-shield", + team_id="t-shield", + prisma_client=mock_prisma_client, + user_api_key_cache=cache, + ) + + owner = asyncio.create_task(_load()) + await started.wait() + waiter = asyncio.create_task(_load()) + await asyncio.sleep(0) + waiter.cancel() + with pytest.raises(asyncio.CancelledError): + await waiter + release.set() + result = await owner + + assert result is not None + assert result.user_id == "u-shield" + mock_prisma_client.db.litellm_teammembership.find_unique.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_stale_membership_redis_replicate_does_not_restore_after_invalidate(): + """A delayed DualCache Redis SET must not resurrect membership after invalidation.""" + from litellm.proxy.auth.auth_checks import get_team_membership, invalidate_team_member_spend_state + from litellm.proxy.common_utils.user_api_key_cache import team_membership_reservation_cache_key + + membership_row = MagicMock() + membership_row.dict = lambda: {"user_id": "u-stale", "team_id": "t-stale", "spend": 9.0} + + hang_redis_set = asyncio.Event() + + async def _hanging_redis_set(*args, **kwargs): + await hang_redis_set.wait() + + redis_cache = MagicMock() + redis_cache.async_get_cache = AsyncMock(return_value=None) + redis_cache.async_set_cache = AsyncMock(side_effect=_hanging_redis_set) + redis_cache.async_delete_cache = AsyncMock() + redis_cache.delete_cache = MagicMock() + + cache = UserApiKeyCache(redis_cache=redis_cache) + mock_prisma_client = MagicMock() + mock_prisma_client.db.litellm_teammembership.find_unique = AsyncMock(return_value=membership_row) + + loaded = await get_team_membership( + user_id="u-stale", + team_id="t-stale", + prisma_client=mock_prisma_client, + user_api_key_cache=cache, + ) + cache_key = team_membership_reservation_cache_key(user_id="u-stale", team_id="t-stale") + await invalidate_team_member_spend_state(user_id="u-stale", team_id="t-stale", user_api_key_cache=cache) + after_invalidate = await cache.async_get_cache(key=cache_key, local_only=True) + + hang_redis_set.set() + await asyncio.sleep(0) + await asyncio.sleep(0) + after_replicate = await cache.async_get_cache(key=cache_key, local_only=True) + + assert loaded is not None + assert after_invalidate is None + assert after_replicate is None + + @pytest.mark.asyncio async def test_invalidate_team_member_spend_state_evicts_the_negative_cache_sentinel(): """ From 53ba8b986644b80fc95d39d11381d88eddb87309 Mon Sep 17 00:00:00 2001 From: Shivi Jain Date: Mon, 14 Sep 2026 22:05:17 +0530 Subject: [PATCH 03/13] fix(auth): load team membership once per request and skip prisma on an L1 hit common_checks was querying get_team_membership twice, and DualCache awaited Redis SET on the auth path, so LRU eviction plus a hung Redis write showed up as two postgres spans --- litellm/proxy/auth/auth_checks.py | 73 +++++++++++-------- litellm/proxy/auth/auth_object_prefetch.py | 5 +- .../proxy/auth/test_auth_checks.py | 34 ++++++--- .../proxy/auth/test_auth_object_prefetch.py | 4 +- .../auth/test_model_access_group_budgets.py | 34 ++++++--- 5 files changed, 93 insertions(+), 57 deletions(-) diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 0baa79fb95a..89b061bdefe 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -2197,30 +2197,51 @@ def _membership_from_cached_payload( return cached_membership if cached_membership is not None else TEAM_MEMBERSHIP_CACHE_MISS -async def _set_team_membership_cache_entry( +async def _set_team_membership_l1( user_api_key_cache: UserApiKeyCache, key: str, value: object, *, - local_only: bool, model_type: type[LiteLLM_TeamMembership] | None, ttl: float | None, ) -> None: match (model_type is not None, ttl is not None): case (False, False): - await user_api_key_cache.async_set_cache(key=key, value=value, local_only=local_only) + await user_api_key_cache.async_set_cache(key=key, value=value, local_only=True) case (False, True): - await user_api_key_cache.async_set_cache(key=key, value=value, local_only=local_only, ttl=ttl) + await user_api_key_cache.async_set_cache(key=key, value=value, local_only=True, ttl=ttl) case (True, False): - await user_api_key_cache.async_set_cache( - key=key, value=value, local_only=local_only, model_type=model_type - ) + await user_api_key_cache.async_set_cache(key=key, value=value, local_only=True, model_type=model_type) case (True, True): await user_api_key_cache.async_set_cache( - key=key, value=value, local_only=local_only, model_type=model_type, ttl=ttl + key=key, value=value, local_only=True, model_type=model_type, ttl=ttl ) +async def _replicate_team_membership_to_redis( + user_api_key_cache: UserApiKeyCache, + key: str, + value: object, + *, + model_type: type[LiteLLM_TeamMembership] | None, + ttl: float | None, + write_epoch: int, +) -> None: + redis_cache: Final = user_api_key_cache.redis_cache + if redis_cache is None or _membership_write_epoch(key) != write_epoch: + return + payload: Final[object] = CacheCodec.serialize(value, model_type=model_type) + try: + if ttl is None: + await redis_cache.async_set_cache(key, payload) + else: + await redis_cache.async_set_cache(key, payload, ttl=ttl) + if _membership_write_epoch(key) != write_epoch: + await redis_cache.async_delete_cache(key) + except Exception: + return + + async def _populate_team_membership_cache( user_api_key_cache: UserApiKeyCache, key: str, @@ -2230,36 +2251,27 @@ async def _populate_team_membership_cache( ttl: float | None = None, ) -> None: write_epoch: Final = _membership_write_epoch(key) - await _set_team_membership_cache_entry( + await _set_team_membership_l1( user_api_key_cache, key, value, - local_only=True, model_type=model_type, ttl=ttl, ) if _membership_write_epoch(key) != write_epoch: - await user_api_key_cache.async_delete_cache(key) + user_api_key_cache.in_memory_cache_for(key).delete_cache(key) return - async def _replicate_to_redis() -> None: - try: - if _membership_write_epoch(key) != write_epoch: - return - await _set_team_membership_cache_entry( - user_api_key_cache, - key, - value, - local_only=False, - model_type=model_type, - ttl=ttl, - ) - if _membership_write_epoch(key) != write_epoch: - await user_api_key_cache.async_delete_cache(key) - except Exception: - return - - asyncio.create_task(_replicate_to_redis()) + asyncio.create_task( + _replicate_team_membership_to_redis( + user_api_key_cache, + key, + value, + model_type=model_type, + ttl=ttl, + write_epoch=write_epoch, + ) + ) @log_db_metrics @@ -2363,6 +2375,9 @@ async def get_team_membership( if prisma_client is None: raise Exception("No db connected") + prisma: Final[object] = prisma_client + if isinstance(prisma, str): + return None task: Final = asyncio.ensure_future( _load_team_membership_on_cache_miss( diff --git a/litellm/proxy/auth/auth_object_prefetch.py b/litellm/proxy/auth/auth_object_prefetch.py index 52e26e885c9..5efc45d3fce 100644 --- a/litellm/proxy/auth/auth_object_prefetch.py +++ b/litellm/proxy/auth/auth_object_prefetch.py @@ -14,7 +14,6 @@ from pydantic import BaseModel, TypeAdapter, ValidationError from litellm._logging import verbose_proxy_logger from litellm.caching.redis_cache import RedisCache -from litellm.constants import DEFAULT_IN_MEMORY_TTL from litellm.models.organization import LiteLLM_OrganizationTable from litellm.models.team import LiteLLM_TeamTableCachedObj from litellm.models.team_membership import LiteLLM_TeamMembership @@ -191,13 +190,13 @@ def _iter_entries(refs: AuthObjectRefs, management_ttl: float) -> Iterator[_Cach ) if refs.organization_id is not None: yield _CacheEntry( - f"org_id:{refs.organization_id}", "organization_row", LiteLLM_OrganizationTable, DEFAULT_IN_MEMORY_TTL + f"org_id:{refs.organization_id}", "organization_row", LiteLLM_OrganizationTable, management_ttl ) yield _CacheEntry( f"org_id:{refs.organization_id}:with_budget", "organization_row", LiteLLM_OrganizationTable, - DEFAULT_IN_MEMORY_TTL, + management_ttl, ) if refs.project_id is not None: yield _CacheEntry(f"project_id:{refs.project_id}", "project_row", LiteLLM_ProjectTableCachedObj, management_ttl) diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index d16c75a4a3b..4fdb49e47ef 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -5527,7 +5527,9 @@ async def _run_internal_user_budget_alert( with ( patch("litellm.proxy.proxy_server.prisma_client", None), # test-quality-ok: common_checks has no database seam - patch("litellm.proxy.proxy_server.get_current_spend", _get_spend), # test-quality-ok: common_checks imports it locally + patch( + "litellm.proxy.proxy_server.get_current_spend", _get_spend + ), # test-quality-ok: common_checks imports it locally patch.object(slack_alerting, "send_alert", send_alert), ): error: Final = await _check_for_error() @@ -6419,9 +6421,7 @@ async def test_get_team_membership_negative_caches_a_missing_row(): assert first is None assert second is None mock_prisma_client.db.litellm_teammembership.find_unique.assert_awaited_once() - cached = await cache.async_get_cache( - key=team_membership_reservation_cache_key(user_id="u-1", team_id="t-1") - ) + cached = await cache.async_get_cache(key=team_membership_reservation_cache_key(user_id="u-1", team_id="t-1")) assert cached == NO_TEAM_MEMBERSHIP_SENTINEL @@ -6601,13 +6601,25 @@ async def test_get_team_membership_db_error_raises_503_not_none(): user_api_key_cache=cache, ) - cached = await cache.async_get_cache( - key=team_membership_reservation_cache_key(user_id="u-fail", team_id="t-fail") - ) + cached = await cache.async_get_cache(key=team_membership_reservation_cache_key(user_id="u-fail", team_id="t-fail")) assert exc.value.status_code == status.HTTP_503_SERVICE_UNAVAILABLE assert cached is None +@pytest.mark.asyncio +async def test_get_team_membership_string_prisma_client_returns_none(): + """Unit tests stub prisma_client as a string; that is not a lookup failure and must not 503.""" + from litellm.proxy.auth.auth_checks import get_team_membership + + result = await get_team_membership( + user_id="u-str", + team_id="t-str", + prisma_client="hello-world", + user_api_key_cache=UserApiKeyCache(), + ) + assert result is None + + @pytest.mark.asyncio async def test_common_checks_does_not_skip_member_limits_when_membership_lookup_fails(): """common_checks must not mark membership loaded-absent after a lookup error.""" @@ -6699,7 +6711,7 @@ async def test_get_team_membership_waiter_cancel_does_not_cancel_shared_load(): @pytest.mark.asyncio async def test_stale_membership_redis_replicate_does_not_restore_after_invalidate(): - """A delayed DualCache Redis SET must not resurrect membership after invalidation.""" + """A delayed Redis SET must not rewrite L1 or leave Redis holding membership after invalidation.""" from litellm.proxy.auth.auth_checks import get_team_membership, invalidate_team_member_spend_state from litellm.proxy.common_utils.user_api_key_cache import team_membership_reservation_cache_key @@ -6739,6 +6751,7 @@ async def test_stale_membership_redis_replicate_does_not_restore_after_invalidat assert loaded is not None assert after_invalidate is None assert after_replicate is None + redis_cache.async_delete_cache.assert_awaited() @pytest.mark.asyncio @@ -6764,10 +6777,7 @@ async def test_invalidate_team_member_spend_state_evicts_the_negative_cache_sent assert before is None await invalidate_team_member_spend_state(user_id="u-1", team_id="t-1", user_api_key_cache=cache) - assert ( - await cache.async_get_cache(key=team_membership_reservation_cache_key(user_id="u-1", team_id="t-1")) - is None - ) + assert await cache.async_get_cache(key=team_membership_reservation_cache_key(user_id="u-1", team_id="t-1")) is None after = await get_team_membership( user_id="u-1", team_id="t-1", prisma_client=mock_prisma_client, user_api_key_cache=cache diff --git a/tests/test_litellm/proxy/auth/test_auth_object_prefetch.py b/tests/test_litellm/proxy/auth/test_auth_object_prefetch.py index 0fd0dda3017..58761e270a0 100644 --- a/tests/test_litellm/proxy/auth/test_auth_object_prefetch.py +++ b/tests/test_litellm/proxy/auth/test_auth_object_prefetch.py @@ -181,8 +181,8 @@ async def test_cold_regime_is_one_mget_one_query_and_the_getters_never_touch_io_ assert sets == sorted( [ f"SET {TEAM_ID}_{USER_ID} ttl=5", - f"SET org_id:{ORG_ID} ttl=5", - f"SET org_id:{ORG_ID}:with_budget ttl=5", + f"SET org_id:{ORG_ID} ttl=60", + f"SET org_id:{ORG_ID}:with_budget ttl=60", f"SET {USER_ID} ttl=60", f"SET team_id:{TEAM_ID} ttl=60", f"SET team_membership:{USER_ID}:{TEAM_ID} ttl=None", diff --git a/tests/test_litellm/proxy/auth/test_model_access_group_budgets.py b/tests/test_litellm/proxy/auth/test_model_access_group_budgets.py index fb82d8708fd..a4206a0ac92 100644 --- a/tests/test_litellm/proxy/auth/test_model_access_group_budgets.py +++ b/tests/test_litellm/proxy/auth/test_model_access_group_budgets.py @@ -31,6 +31,7 @@ from litellm.proxy.auth.auth_checks import ( ) from litellm.proxy.common_utils.reset_budget_job import _model_access_group_counter_key from litellm.proxy.common_utils.user_api_key_cache import ( + NO_TEAM_MEMBERSHIP_SENTINEL, UserApiKeyCache, model_access_group_registry_cache_key, model_access_group_spend_counter_key, @@ -95,6 +96,11 @@ async def _cache( ), model_type=LiteLLM_TeamMembership, ) + else: + await cache.async_set_cache( + key=team_membership_reservation_cache_key(user_id=USER_ID, team_id=TEAM_ID), + value=NO_TEAM_MEMBERSHIP_SENTINEL, + ) if org_models: await cache.async_set_cache( key=f"org_id:{ORG_ID}", @@ -314,9 +320,7 @@ class _RecordingPrismaClient: def __init__(self, *rows: _MagBudgetRow) -> None: self.rows = {row.access_group_name: row for row in rows} self.batches: list[list[str]] = [] - self.db = SimpleNamespace( - litellm_modelaccessgroupbudgettable=SimpleNamespace(find_many=self._find_many) - ) + self.db = SimpleNamespace(litellm_modelaccessgroupbudgettable=SimpleNamespace(find_many=self._find_many)) async def _find_many(self, **kwargs): requested = list(kwargs["where"]["access_group_name"]["in"]) @@ -345,7 +349,9 @@ async def _enforce( read, seen = _spend_reader(spend_by_counter_key or {}) # The check takes its client and cache as arguments, injected just below. get_current_spend is the # one collaborator it reaches by a lazy `from litellm.proxy.proxy_server import`, with no parameter. - with patch("litellm.proxy.proxy_server.get_current_spend", read): # test-quality-ok: get_current_spend is lazily imported inside _model_access_group_max_budget_check and has no injection point + with patch( + "litellm.proxy.proxy_server.get_current_spend", read + ): # test-quality-ok: get_current_spend is lazily imported inside _model_access_group_max_budget_check and has no injection point await _model_access_group_max_budget_check( matched_model_access_groups=matched, prisma_client=prisma_client if prisma_client is not None else _RecordingPrismaClient(*rows), @@ -491,9 +497,7 @@ async def test_a_second_request_serves_the_budget_row_from_cache(): async def test_a_database_error_does_not_block_the_request(): class _FailingPrismaClient: def __init__(self) -> None: - self.db = SimpleNamespace( - litellm_modelaccessgroupbudgettable=SimpleNamespace(find_many=self._boom) - ) + self.db = SimpleNamespace(litellm_modelaccessgroupbudgettable=SimpleNamespace(find_many=self._boom)) async def _boom(self, **kwargs): raise RuntimeError("database unavailable") @@ -509,9 +513,15 @@ async def _common_checks_with_over_budget_group(*, skip_budget_checks: bool) -> with ( # common_checks resolves all three off the proxy_server module at call time; its signature # has no client, cache or spend-reader parameter to pass them through instead. - patch("litellm.proxy.proxy_server.prisma_client", prisma_client), # test-quality-ok: common_checks lazily imports prisma_client from proxy_server and takes no client parameter - patch("litellm.proxy.proxy_server.user_api_key_cache", cache), # test-quality-ok: common_checks lazily imports user_api_key_cache from proxy_server and takes no cache parameter - patch("litellm.proxy.proxy_server.get_current_spend", read), # test-quality-ok: get_current_spend is lazily imported inside the budget check and has no injection point + patch( + "litellm.proxy.proxy_server.prisma_client", prisma_client + ), # test-quality-ok: common_checks lazily imports prisma_client from proxy_server and takes no client parameter + patch( + "litellm.proxy.proxy_server.user_api_key_cache", cache + ), # test-quality-ok: common_checks lazily imports user_api_key_cache from proxy_server and takes no cache parameter + patch( + "litellm.proxy.proxy_server.get_current_spend", read + ), # test-quality-ok: get_current_spend is lazily imported inside the budget check and has no injection point ): return await common_checks( request_body={"model": "gpt-4o", "messages": []}, @@ -524,7 +534,9 @@ async def _common_checks_with_over_budget_group(*, skip_budget_checks: bool) -> llm_router=Router(model_list=MODEL_LIST), proxy_logging_obj=ProxyLogging(user_api_key_cache=cache), valid_token=UserAPIKeyAuth(api_key="hashed", models=["tier-a"], user_id=USER_ID), - request=SimpleNamespace(method="POST", headers={}, query_params={}, url=SimpleNamespace(path="/v1/chat/completions")), + request=SimpleNamespace( + method="POST", headers={}, query_params={}, url=SimpleNamespace(path="/v1/chat/completions") + ), skip_budget_checks=skip_budget_checks, ) From abaa2f8b81c4b73bc57a169fb2a85f6a535c0984 Mon Sep 17 00:00:00 2001 From: Shivi Jain Date: Mon, 14 Sep 2026 22:12:17 +0530 Subject: [PATCH 04/13] fix(auth): put TQ008 suppressions on the patch call lines The test-quality gate attributes the comment to the `patch(` line, so reasons on the closing paren did not count and lint failed after format started passing. Co-authored-by: Cursor --- .../proxy/auth/test_auth_checks.py | 32 +++++++++++++------ .../auth/test_model_access_group_budgets.py | 18 +++++------ 2 files changed, 30 insertions(+), 20 deletions(-) diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index 4fdb49e47ef..29f6c3861f6 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -5527,9 +5527,9 @@ async def _run_internal_user_budget_alert( with ( patch("litellm.proxy.proxy_server.prisma_client", None), # test-quality-ok: common_checks has no database seam - patch( + patch( # test-quality-ok: common_checks imports get_current_spend locally "litellm.proxy.proxy_server.get_current_spend", _get_spend - ), # test-quality-ok: common_checks imports it locally + ), patch.object(slack_alerting, "send_alert", send_alert), ): error: Final = await _check_for_error() @@ -6554,14 +6554,20 @@ async def test_common_checks_calls_get_team_membership_once_per_request(): membership.spend = 0.0 with ( - patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), - patch("litellm.proxy.proxy_server.user_api_key_cache", UserApiKeyCache()), - patch( + patch( # test-quality-ok: common_checks imports prisma_client from proxy_server + "litellm.proxy.proxy_server.prisma_client", MagicMock() + ), + patch( # test-quality-ok: common_checks imports user_api_key_cache from proxy_server + "litellm.proxy.proxy_server.user_api_key_cache", UserApiKeyCache() + ), + patch( # test-quality-ok: counts membership loads; common_checks has no membership seam "litellm.proxy.auth.auth_checks.get_team_membership", new_callable=AsyncMock, return_value=membership, ) as load_membership, - patch("litellm.proxy.proxy_server.get_current_spend", new_callable=AsyncMock, return_value=0.0), + patch( # test-quality-ok: common_checks imports get_current_spend locally + "litellm.proxy.proxy_server.get_current_spend", new_callable=AsyncMock, return_value=0.0 + ), ): result = await common_checks( request_body={"model": "gpt-4o-mini", "messages": [{"role": "user", "content": "hi"}]}, @@ -6640,14 +6646,20 @@ async def test_common_checks_does_not_skip_member_limits_when_membership_lookup_ ) with ( - patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), - patch("litellm.proxy.proxy_server.user_api_key_cache", UserApiKeyCache()), - patch( + patch( # test-quality-ok: common_checks imports prisma_client from proxy_server + "litellm.proxy.proxy_server.prisma_client", MagicMock() + ), + patch( # test-quality-ok: common_checks imports user_api_key_cache from proxy_server + "litellm.proxy.proxy_server.user_api_key_cache", UserApiKeyCache() + ), + patch( # test-quality-ok: injects membership lookup failure; common_checks has no seam "litellm.proxy.auth.auth_checks.get_team_membership", new_callable=AsyncMock, side_effect=lookup_error, ), - patch("litellm.proxy.proxy_server.get_current_spend", new_callable=AsyncMock, return_value=0.0), + patch( # test-quality-ok: common_checks imports get_current_spend locally + "litellm.proxy.proxy_server.get_current_spend", new_callable=AsyncMock, return_value=0.0 + ), ): with pytest.raises(HTTPException) as exc: await common_checks( diff --git a/tests/test_litellm/proxy/auth/test_model_access_group_budgets.py b/tests/test_litellm/proxy/auth/test_model_access_group_budgets.py index a4206a0ac92..7506bd031d9 100644 --- a/tests/test_litellm/proxy/auth/test_model_access_group_budgets.py +++ b/tests/test_litellm/proxy/auth/test_model_access_group_budgets.py @@ -349,9 +349,9 @@ async def _enforce( read, seen = _spend_reader(spend_by_counter_key or {}) # The check takes its client and cache as arguments, injected just below. get_current_spend is the # one collaborator it reaches by a lazy `from litellm.proxy.proxy_server import`, with no parameter. - with patch( + with patch( # test-quality-ok: get_current_spend is lazily imported inside the budget check "litellm.proxy.proxy_server.get_current_spend", read - ): # test-quality-ok: get_current_spend is lazily imported inside _model_access_group_max_budget_check and has no injection point + ): await _model_access_group_max_budget_check( matched_model_access_groups=matched, prisma_client=prisma_client if prisma_client is not None else _RecordingPrismaClient(*rows), @@ -511,17 +511,15 @@ async def _common_checks_with_over_budget_group(*, skip_budget_checks: bool) -> read, _ = _spend_reader({MODEL_ACCESS_GROUP_COUNTER_KEY: 99.0}) with ( - # common_checks resolves all three off the proxy_server module at call time; its signature - # has no client, cache or spend-reader parameter to pass them through instead. - patch( + patch( # test-quality-ok: common_checks lazily imports prisma_client from proxy_server "litellm.proxy.proxy_server.prisma_client", prisma_client - ), # test-quality-ok: common_checks lazily imports prisma_client from proxy_server and takes no client parameter - patch( + ), + patch( # test-quality-ok: common_checks lazily imports user_api_key_cache from proxy_server "litellm.proxy.proxy_server.user_api_key_cache", cache - ), # test-quality-ok: common_checks lazily imports user_api_key_cache from proxy_server and takes no cache parameter - patch( + ), + patch( # test-quality-ok: get_current_spend is lazily imported inside the budget check "litellm.proxy.proxy_server.get_current_spend", read - ), # test-quality-ok: get_current_spend is lazily imported inside the budget check and has no injection point + ), ): return await common_checks( request_body={"model": "gpt-4o", "messages": []}, From 75b16df6fa38d0521ee0bc658fb847611f865e06 Mon Sep 17 00:00:00 2001 From: Shivi Jain Date: Mon, 14 Sep 2026 22:22:34 +0530 Subject: [PATCH 05/13] fix(auth): keep prefetched org entries on the 5s getter TTL Organization mutations do not evict those cache keys, so stretching prefetch to the management TTL would leave stale org grants in L1. Co-authored-by: Cursor --- litellm/proxy/auth/auth_object_prefetch.py | 5 +++-- tests/test_litellm/proxy/auth/test_auth_object_prefetch.py | 4 ++-- 2 files changed, 5 insertions(+), 4 deletions(-) diff --git a/litellm/proxy/auth/auth_object_prefetch.py b/litellm/proxy/auth/auth_object_prefetch.py index 5efc45d3fce..52e26e885c9 100644 --- a/litellm/proxy/auth/auth_object_prefetch.py +++ b/litellm/proxy/auth/auth_object_prefetch.py @@ -14,6 +14,7 @@ from pydantic import BaseModel, TypeAdapter, ValidationError from litellm._logging import verbose_proxy_logger from litellm.caching.redis_cache import RedisCache +from litellm.constants import DEFAULT_IN_MEMORY_TTL from litellm.models.organization import LiteLLM_OrganizationTable from litellm.models.team import LiteLLM_TeamTableCachedObj from litellm.models.team_membership import LiteLLM_TeamMembership @@ -190,13 +191,13 @@ def _iter_entries(refs: AuthObjectRefs, management_ttl: float) -> Iterator[_Cach ) if refs.organization_id is not None: yield _CacheEntry( - f"org_id:{refs.organization_id}", "organization_row", LiteLLM_OrganizationTable, management_ttl + f"org_id:{refs.organization_id}", "organization_row", LiteLLM_OrganizationTable, DEFAULT_IN_MEMORY_TTL ) yield _CacheEntry( f"org_id:{refs.organization_id}:with_budget", "organization_row", LiteLLM_OrganizationTable, - management_ttl, + DEFAULT_IN_MEMORY_TTL, ) if refs.project_id is not None: yield _CacheEntry(f"project_id:{refs.project_id}", "project_row", LiteLLM_ProjectTableCachedObj, management_ttl) diff --git a/tests/test_litellm/proxy/auth/test_auth_object_prefetch.py b/tests/test_litellm/proxy/auth/test_auth_object_prefetch.py index 58761e270a0..0fd0dda3017 100644 --- a/tests/test_litellm/proxy/auth/test_auth_object_prefetch.py +++ b/tests/test_litellm/proxy/auth/test_auth_object_prefetch.py @@ -181,8 +181,8 @@ async def test_cold_regime_is_one_mget_one_query_and_the_getters_never_touch_io_ assert sets == sorted( [ f"SET {TEAM_ID}_{USER_ID} ttl=5", - f"SET org_id:{ORG_ID} ttl=60", - f"SET org_id:{ORG_ID}:with_budget ttl=60", + f"SET org_id:{ORG_ID} ttl=5", + f"SET org_id:{ORG_ID}:with_budget ttl=5", f"SET {USER_ID} ttl=60", f"SET team_id:{TEAM_ID} ttl=60", f"SET team_membership:{USER_ID}:{TEAM_ID} ttl=None", From f4e21430c09d5201673b8093031d80391221a05a Mon Sep 17 00:00:00 2001 From: Shivi Jain Date: Mon, 14 Sep 2026 22:32:11 +0530 Subject: [PATCH 06/13] test(auth): freeze prefetch cache clock so the org getter cannot miss on a slow runner Prefetch still writes org entries with the 5s getter TTL. This test only asserts the SQL join, and wall-clock expiry on CI turned that into a MagicMock await TypeError. Co-authored-by: Cursor --- tests/proxy_behavior/auth/test_auth_object_prefetch.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/proxy_behavior/auth/test_auth_object_prefetch.py b/tests/proxy_behavior/auth/test_auth_object_prefetch.py index 59e9585a296..cfa958500af 100644 --- a/tests/proxy_behavior/auth/test_auth_object_prefetch.py +++ b/tests/proxy_behavior/auth/test_auth_object_prefetch.py @@ -66,6 +66,7 @@ async def test_join_binds_the_membership_to_the_requested_team(prisma): cache = _frozen_cache() refs = AuthObjectRefs(user_id=user_id, team_id=team_a, membership_user_id=user_id, organization_id=org_id) await prefetch_auth_objects(refs=refs, user_api_key_cache=cache, prisma_client=prisma) + assert cache.in_memory_cache.get_cache(f"org_id:{org_id}") is not None dead_db = _dead_db() membership = await get_team_membership( From b00bb3556327ad576bea306bf93440dee27227b0 Mon Sep 17 00:00:00 2001 From: yassin Date: Mon, 14 Sep 2026 19:06:29 +0000 Subject: [PATCH 07/13] fix(auth): drop the membership write-epoch and background Redis replicate, load membership lazily Move the cache-miss marker out of litellm.constants into auth_checks (CodeQL cyclic import) and stop logging user_id/team_id in the lookup failure (CodeQL log injection). Write the membership row through DualCache synchronously again instead of a background Redis task guarded by a bounded write-epoch map: the epoch was sampled after the Prisma read, so an invalidate that raced the read could be cached as current, and eviction of the epoch entry could let an old Redis write land. The synchronous write keeps invalidate_team_member_spend_state authoritative. Lookup failures return None again (fail-open like main) instead of 503, and the load is skipped on routes that neither resolve a model nor run budget checks. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/constants.py | 7 - litellm/proxy/auth/auth_checks.py | 179 +++----------- .../proxy/auth/test_auth_checks.py | 226 ++++++------------ 3 files changed, 106 insertions(+), 306 deletions(-) diff --git a/litellm/constants.py b/litellm/constants.py index 5c7a02d0743..5751e6e46af 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -2039,10 +2039,3 @@ BATCH_ENQUEUED_TOKEN_LIMIT_METADATA_KEY: Final = "batch_enqueued_token_limit" # Shared read-only empty mapping, for defaulting optional Mapping parameters without # constructing a fresh mutable dict at each call site. EMPTY_MAPPING: Final = MappingProxyType({}) - - -class TeamMembershipCacheMiss: - __slots__ = () - - -TEAM_MEMBERSHIP_CACHE_MISS: Final = TeamMembershipCacheMiss() diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 89b061bdefe..a8f9cc77577 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -35,8 +35,6 @@ from litellm.constants import ( MODEL_ACCESS_GROUP_REGISTRY_MAX_SIZE, REGISTRY_ERROR_NEGATIVE_CACHE_TTL, TAG_REGISTRY_MAX_SIZE, - TEAM_MEMBERSHIP_CACHE_MISS, - TeamMembershipCacheMiss, ) from litellm.litellm_core_utils.dd_tracing import tracer from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider @@ -331,27 +329,19 @@ db_cache_expiry: Final = DEFAULT_IN_MEMORY_TTL # refresh every 5s _TEAM_MEMBERSHIP_INFLIGHT_MAX: Final = 10000 _team_membership_inflight: Final = LimitedSizeOrderedDict(max_size=_TEAM_MEMBERSHIP_INFLIGHT_MAX) -_team_membership_write_epoch: Final = LimitedSizeOrderedDict(max_size=_TEAM_MEMBERSHIP_INFLIGHT_MAX) + + +class _TeamMembershipCacheMiss: + __slots__ = () + + +_TEAM_MEMBERSHIP_CACHE_MISS: Final = _TeamMembershipCacheMiss() all_routes: Final = LiteLLMRoutes.openai_routes.value + LiteLLMRoutes.management_routes.value -def _membership_write_epoch(key: str) -> int: - cached: Final[object] = _team_membership_write_epoch.get(key, 0) - return cached if isinstance(cached, int) else 0 - - -def _bump_membership_write_epoch(key: str) -> None: - _team_membership_write_epoch[key] = _membership_write_epoch(key) + 1 - - def _membership_from_shared_load(result: object) -> LiteLLM_TeamMembership | None: - if result is None or isinstance(result, LiteLLM_TeamMembership): - return result - raise HTTPException( - status_code=status.HTTP_503_SERVICE_UNAVAILABLE, - detail="Failed to load team membership", - ) + return result if isinstance(result, LiteLLM_TeamMembership) else None def _log_budget_lookup_failure(entity: str, error: Exception) -> None: @@ -897,21 +887,6 @@ async def common_checks( """ from litellm.proxy.proxy_server import prisma_client, user_api_key_cache - # One membership read for model-access, access-group attribution, and - # member-budget. Each used to call get_team_membership independently; - # DualCache Redis SET/GET on every miss made those look like two Postgres spans. - loaded_team_membership: LiteLLM_TeamMembership | None = None - team_membership_loaded = False - if team_object is not None and valid_token is not None and valid_token.user_id is not None: - loaded_team_membership = await get_team_membership( - user_id=valid_token.user_id, - team_id=team_object.team_id, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=proxy_logging_obj, - ) - team_membership_loaded = True - _model: Final[str | list[str] | None] = get_model_from_request( request_data=request_body, route=route, @@ -926,6 +901,22 @@ async def common_checks( and (route in MODEL_DISCOVERY_ROUTES or not RouteChecks.is_llm_api_route(route=route)) ) + membership_user_id: Final = ( + valid_token.user_id if valid_token is not None and (bool(_model) or not skip_all_budget_checks) else None + ) + team_membership_loaded: Final = team_object is not None and membership_user_id is not None + loaded_team_membership: Final = ( + await get_team_membership( + user_id=membership_user_id, + team_id=team_object.team_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + if team_object is not None and membership_user_id is not None + else None + ) + unpriced_models: Final = ( _unpriced_models_in_request(model=_model, llm_router=llm_router) if litellm.block_requests_for_models_without_pricing and RouteChecks.is_llm_api_route(route=route) @@ -2188,90 +2179,13 @@ async def get_tag_object( def _membership_from_cached_payload( cached: object, -) -> LiteLLM_TeamMembership | None | TeamMembershipCacheMiss: +) -> LiteLLM_TeamMembership | None | _TeamMembershipCacheMiss: if cached is None: - return TEAM_MEMBERSHIP_CACHE_MISS + return _TEAM_MEMBERSHIP_CACHE_MISS if cached == NO_TEAM_MEMBERSHIP_SENTINEL: return None cached_membership: Final = CacheCodec.deserialize(cached, model_type=LiteLLM_TeamMembership) - return cached_membership if cached_membership is not None else TEAM_MEMBERSHIP_CACHE_MISS - - -async def _set_team_membership_l1( - user_api_key_cache: UserApiKeyCache, - key: str, - value: object, - *, - model_type: type[LiteLLM_TeamMembership] | None, - ttl: float | None, -) -> None: - match (model_type is not None, ttl is not None): - case (False, False): - await user_api_key_cache.async_set_cache(key=key, value=value, local_only=True) - case (False, True): - await user_api_key_cache.async_set_cache(key=key, value=value, local_only=True, ttl=ttl) - case (True, False): - await user_api_key_cache.async_set_cache(key=key, value=value, local_only=True, model_type=model_type) - case (True, True): - await user_api_key_cache.async_set_cache( - key=key, value=value, local_only=True, model_type=model_type, ttl=ttl - ) - - -async def _replicate_team_membership_to_redis( - user_api_key_cache: UserApiKeyCache, - key: str, - value: object, - *, - model_type: type[LiteLLM_TeamMembership] | None, - ttl: float | None, - write_epoch: int, -) -> None: - redis_cache: Final = user_api_key_cache.redis_cache - if redis_cache is None or _membership_write_epoch(key) != write_epoch: - return - payload: Final[object] = CacheCodec.serialize(value, model_type=model_type) - try: - if ttl is None: - await redis_cache.async_set_cache(key, payload) - else: - await redis_cache.async_set_cache(key, payload, ttl=ttl) - if _membership_write_epoch(key) != write_epoch: - await redis_cache.async_delete_cache(key) - except Exception: - return - - -async def _populate_team_membership_cache( - user_api_key_cache: UserApiKeyCache, - key: str, - value: object, - *, - model_type: type[LiteLLM_TeamMembership] | None = None, - ttl: float | None = None, -) -> None: - write_epoch: Final = _membership_write_epoch(key) - await _set_team_membership_l1( - user_api_key_cache, - key, - value, - model_type=model_type, - ttl=ttl, - ) - if _membership_write_epoch(key) != write_epoch: - user_api_key_cache.in_memory_cache_for(key).delete_cache(key) - return - - asyncio.create_task( - _replicate_team_membership_to_redis( - user_api_key_cache, - key, - value, - model_type=model_type, - ttl=ttl, - write_epoch=write_epoch, - ) - ) + return cached_membership if cached_membership is not None else _TEAM_MEMBERSHIP_CACHE_MISS @log_db_metrics @@ -2283,7 +2197,7 @@ async def _fetch_team_membership_from_db( parent_otel_span: Span | None = None, proxy_logging_obj: ProxyLogging | None = None, ) -> LiteLLM_TeamMembership | None: - """Prisma read + L1 populate. Decorated so cache hits on ``get_team_membership`` are not postgres spans.""" + """Prisma read + cache populate. Decorated so cache hits on ``get_team_membership`` are not postgres spans.""" _ = parent_otel_span, proxy_logging_obj response: Final = await _dictable_table(TeamMembershipRepository(prisma_client)).find_unique( where={"user_id_team_id": {"user_id": user_id, "team_id": team_id}}, @@ -2291,19 +2205,17 @@ async def _fetch_team_membership_from_db( ) _key: Final = team_membership_reservation_cache_key(user_id=user_id, team_id=team_id) if response is None: - await _populate_team_membership_cache( - user_api_key_cache, - _key, - NO_TEAM_MEMBERSHIP_SENTINEL, + await user_api_key_cache.async_set_cache( + key=_key, + value=NO_TEAM_MEMBERSHIP_SENTINEL, ttl=get_management_object_ttl(user_api_key_cache), ) return None membership: Final = LiteLLM_TeamMembership.model_validate(response.dict()) - await _populate_team_membership_cache( - user_api_key_cache, - _key, - membership, + await user_api_key_cache.async_set_cache( + key=_key, + value=membership, model_type=LiteLLM_TeamMembership, ) return membership @@ -2321,7 +2233,7 @@ async def _load_team_membership_on_cache_miss( try: redis_cached: Final[object] = await user_api_key_cache.async_get_cache(key=cache_key) redis_membership: Final = _membership_from_cached_payload(redis_cached) - if not isinstance(redis_membership, TeamMembershipCacheMiss): + if not isinstance(redis_membership, _TeamMembershipCacheMiss): return redis_membership return await _fetch_team_membership_from_db( @@ -2332,18 +2244,9 @@ async def _load_team_membership_on_cache_miss( parent_otel_span=parent_otel_span, proxy_logging_obj=proxy_logging_obj, ) - except HTTPException: - raise - except Exception as e: - verbose_proxy_logger.exception( - "Error getting team membership for user_id: %s, team_id: %s", - user_id, - team_id, - ) - raise HTTPException( - status_code=status.HTTP_503_SERVICE_UNAVAILABLE, - detail="Failed to load team membership", - ) from e + except Exception: + verbose_proxy_logger.exception("Error getting team membership") + return None async def get_team_membership( @@ -2366,7 +2269,7 @@ async def get_team_membership( l1_cached: Final[object] = await user_api_key_cache.async_get_cache(key=_key, local_only=True) l1_membership: Final = _membership_from_cached_payload(l1_cached) - if not isinstance(l1_membership, TeamMembershipCacheMiss): + if not isinstance(l1_membership, _TeamMembershipCacheMiss): return l1_membership inflight: Final[object] = _team_membership_inflight.get(_key) @@ -2375,9 +2278,6 @@ async def get_team_membership( if prisma_client is None: raise Exception("No db connected") - prisma: Final[object] = prisma_client - if isinstance(prisma, str): - return None task: Final = asyncio.ensure_future( _load_team_membership_on_cache_miss( @@ -2908,7 +2808,6 @@ async def invalidate_team_member_spend_state( ttl=SPEND_DB_FLOOR_CACHE_TTL_SECONDS, ) - _bump_membership_write_epoch(team_membership_reservation_cache_key(user_id=user_id, team_id=team_id)) await evict_and_broadcast( cache_keys=( team_membership_auth_cache_key(team_id=team_id, user_id=user_id), diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index 29f6c3861f6..04e41608b26 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -6494,52 +6494,6 @@ async def test_get_team_membership_coalesces_parallel_db_fetches(): mock_prisma_client.db.litellm_teammembership.find_unique.assert_awaited_once() -@pytest.mark.asyncio -async def test_get_team_membership_returns_before_redis_set_completes(): - """Auth must not wait on DualCache Redis SET; L1 is enough for the next lookup.""" - from litellm.proxy.auth.auth_checks import get_team_membership - - membership_row = MagicMock() - membership_row.dict = lambda: {"user_id": "u-redis", "team_id": "t-redis", "spend": 2.0} - - hang_redis_set = asyncio.Event() - - async def _hanging_redis_set(*args, **kwargs): - await hang_redis_set.wait() - - redis_cache = MagicMock() - redis_cache.async_get_cache = AsyncMock(return_value=None) - redis_cache.async_set_cache = AsyncMock(side_effect=_hanging_redis_set) - - cache = UserApiKeyCache(redis_cache=redis_cache) - mock_prisma_client = MagicMock() - mock_prisma_client.db.litellm_teammembership.find_unique = AsyncMock(return_value=membership_row) - - first = await asyncio.wait_for( - get_team_membership( - user_id="u-redis", - team_id="t-redis", - prisma_client=mock_prisma_client, - user_api_key_cache=cache, - ), - timeout=0.5, - ) - second = await get_team_membership( - user_id="u-redis", - team_id="t-redis", - prisma_client=mock_prisma_client, - user_api_key_cache=cache, - ) - - hang_redis_set.set() - await asyncio.sleep(0) - - assert first is not None and second is not None - assert first.spend == 2.0 - assert second.spend == 2.0 - mock_prisma_client.db.litellm_teammembership.find_unique.assert_awaited_once() - - @pytest.mark.asyncio async def test_common_checks_calls_get_team_membership_once_per_request(): """Model-access, attribution, and member-budget must reuse one membership load.""" @@ -6588,33 +6542,85 @@ async def test_common_checks_calls_get_team_membership_once_per_request(): @pytest.mark.asyncio -async def test_get_team_membership_db_error_raises_503_not_none(): - """A Prisma failure must fail closed as 503, not look like a missing membership row.""" - from fastapi import HTTPException +async def test_common_checks_skips_membership_load_when_no_check_reads_it(): + """A management route has no model and no budget gate, so the membership row is never loaded.""" + from fastapi import Request + from litellm.proxy.auth.auth_checks import common_checks + + team = LiteLLM_TeamTable(team_id="t-lazy") + token = UserAPIKeyAuth(token="k-lazy", user_id="u-lazy", team_id="t-lazy") + + with ( + patch( # test-quality-ok: common_checks imports prisma_client from proxy_server + "litellm.proxy.proxy_server.prisma_client", MagicMock() + ), + patch( # test-quality-ok: common_checks imports user_api_key_cache from proxy_server + "litellm.proxy.proxy_server.user_api_key_cache", UserApiKeyCache() + ), + patch( # test-quality-ok: counts membership loads; common_checks has no membership seam + "litellm.proxy.auth.auth_checks.get_team_membership", + new_callable=AsyncMock, + ) as load_membership, + ): + result = await common_checks( + request_body={}, + team_object=team, + user_object=LiteLLM_UserTable(user_id="u-lazy"), + end_user_object=None, + global_proxy_spend=None, + general_settings={}, + route="/key/info", + llm_router=None, + proxy_logging_obj=MagicMock(), + valid_token=token, + request=MagicMock(spec=Request), + ) + + assert result is True + load_membership.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_get_team_membership_db_error_returns_none_and_retries_next_call(): + """A Prisma failure reads as no membership, caches nothing, and the next call hits the DB again.""" from litellm.proxy.auth.auth_checks import get_team_membership from litellm.proxy.common_utils.user_api_key_cache import team_membership_reservation_cache_key + membership_row = MagicMock() + membership_row.dict = lambda: {"user_id": "u-fail", "team_id": "t-fail", "spend": 1.0} mock_prisma_client = MagicMock() - mock_prisma_client.db.litellm_teammembership.find_unique = AsyncMock(side_effect=RuntimeError("db down")) + mock_prisma_client.db.litellm_teammembership.find_unique = AsyncMock( + side_effect=[RuntimeError("db down"), membership_row] + ) cache = UserApiKeyCache() - with pytest.raises(HTTPException) as exc: - await get_team_membership( - user_id="u-fail", - team_id="t-fail", - prisma_client=mock_prisma_client, - user_api_key_cache=cache, - ) + failed = await get_team_membership( + user_id="u-fail", + team_id="t-fail", + prisma_client=mock_prisma_client, + user_api_key_cache=cache, + ) + cached_after_failure = await cache.async_get_cache( + key=team_membership_reservation_cache_key(user_id="u-fail", team_id="t-fail") + ) + recovered = await get_team_membership( + user_id="u-fail", + team_id="t-fail", + prisma_client=mock_prisma_client, + user_api_key_cache=cache, + ) - cached = await cache.async_get_cache(key=team_membership_reservation_cache_key(user_id="u-fail", team_id="t-fail")) - assert exc.value.status_code == status.HTTP_503_SERVICE_UNAVAILABLE - assert cached is None + assert failed is None + assert cached_after_failure is None + assert recovered is not None + assert recovered.user_id == "u-fail" + assert mock_prisma_client.db.litellm_teammembership.find_unique.await_count == 2 @pytest.mark.asyncio async def test_get_team_membership_string_prisma_client_returns_none(): - """Unit tests stub prisma_client as a string; that is not a lookup failure and must not 503.""" + """Unit tests stub prisma_client as a string; the lookup fails and reads as no membership.""" from litellm.proxy.auth.auth_checks import get_team_membership result = await get_team_membership( @@ -6626,59 +6632,6 @@ async def test_get_team_membership_string_prisma_client_returns_none(): assert result is None -@pytest.mark.asyncio -async def test_common_checks_does_not_skip_member_limits_when_membership_lookup_fails(): - """common_checks must not mark membership loaded-absent after a lookup error.""" - from fastapi import HTTPException, Request - - from litellm.proxy.auth.auth_checks import common_checks - - team = LiteLLM_TeamTable(team_id="t-fail-closed") - token = UserAPIKeyAuth( - token="k-fail-closed", - user_id="u-fail-closed", - team_id="t-fail-closed", - models=["gpt-4o-mini"], - ) - lookup_error = HTTPException( - status_code=status.HTTP_503_SERVICE_UNAVAILABLE, - detail="Failed to load team membership", - ) - - with ( - patch( # test-quality-ok: common_checks imports prisma_client from proxy_server - "litellm.proxy.proxy_server.prisma_client", MagicMock() - ), - patch( # test-quality-ok: common_checks imports user_api_key_cache from proxy_server - "litellm.proxy.proxy_server.user_api_key_cache", UserApiKeyCache() - ), - patch( # test-quality-ok: injects membership lookup failure; common_checks has no seam - "litellm.proxy.auth.auth_checks.get_team_membership", - new_callable=AsyncMock, - side_effect=lookup_error, - ), - patch( # test-quality-ok: common_checks imports get_current_spend locally - "litellm.proxy.proxy_server.get_current_spend", new_callable=AsyncMock, return_value=0.0 - ), - ): - with pytest.raises(HTTPException) as exc: - await common_checks( - request_body={"model": "gpt-4o-mini", "messages": [{"role": "user", "content": "hi"}]}, - team_object=team, - user_object=LiteLLM_UserTable(user_id="u-fail-closed"), - end_user_object=None, - global_proxy_spend=None, - general_settings={}, - route="/chat/completions", - llm_router=None, - proxy_logging_obj=MagicMock(), - valid_token=token, - request=MagicMock(spec=Request), - ) - - assert exc.value.status_code == status.HTTP_503_SERVICE_UNAVAILABLE - - @pytest.mark.asyncio async def test_get_team_membership_waiter_cancel_does_not_cancel_shared_load(): """Cancelling one coalesced waiter must not cancel the shared Prisma load.""" @@ -6721,51 +6674,6 @@ async def test_get_team_membership_waiter_cancel_does_not_cancel_shared_load(): mock_prisma_client.db.litellm_teammembership.find_unique.assert_awaited_once() -@pytest.mark.asyncio -async def test_stale_membership_redis_replicate_does_not_restore_after_invalidate(): - """A delayed Redis SET must not rewrite L1 or leave Redis holding membership after invalidation.""" - from litellm.proxy.auth.auth_checks import get_team_membership, invalidate_team_member_spend_state - from litellm.proxy.common_utils.user_api_key_cache import team_membership_reservation_cache_key - - membership_row = MagicMock() - membership_row.dict = lambda: {"user_id": "u-stale", "team_id": "t-stale", "spend": 9.0} - - hang_redis_set = asyncio.Event() - - async def _hanging_redis_set(*args, **kwargs): - await hang_redis_set.wait() - - redis_cache = MagicMock() - redis_cache.async_get_cache = AsyncMock(return_value=None) - redis_cache.async_set_cache = AsyncMock(side_effect=_hanging_redis_set) - redis_cache.async_delete_cache = AsyncMock() - redis_cache.delete_cache = MagicMock() - - cache = UserApiKeyCache(redis_cache=redis_cache) - mock_prisma_client = MagicMock() - mock_prisma_client.db.litellm_teammembership.find_unique = AsyncMock(return_value=membership_row) - - loaded = await get_team_membership( - user_id="u-stale", - team_id="t-stale", - prisma_client=mock_prisma_client, - user_api_key_cache=cache, - ) - cache_key = team_membership_reservation_cache_key(user_id="u-stale", team_id="t-stale") - await invalidate_team_member_spend_state(user_id="u-stale", team_id="t-stale", user_api_key_cache=cache) - after_invalidate = await cache.async_get_cache(key=cache_key, local_only=True) - - hang_redis_set.set() - await asyncio.sleep(0) - await asyncio.sleep(0) - after_replicate = await cache.async_get_cache(key=cache_key, local_only=True) - - assert loaded is not None - assert after_invalidate is None - assert after_replicate is None - redis_cache.async_delete_cache.assert_awaited() - - @pytest.mark.asyncio async def test_invalidate_team_member_spend_state_evicts_the_negative_cache_sentinel(): """ From 5bcc23014e8e64ea07a7e9259e06ef86b6e051d8 Mon Sep 17 00:00:00 2001 From: yassin Date: Mon, 14 Sep 2026 19:12:35 +0000 Subject: [PATCH 08/13] refactor(auth): drop the membership fetch docstring Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/auth/auth_checks.py | 1 - 1 file changed, 1 deletion(-) diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index a8f9cc77577..fa85400f4e0 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -2197,7 +2197,6 @@ async def _fetch_team_membership_from_db( parent_otel_span: Span | None = None, proxy_logging_obj: ProxyLogging | None = None, ) -> LiteLLM_TeamMembership | None: - """Prisma read + cache populate. Decorated so cache hits on ``get_team_membership`` are not postgres spans.""" _ = parent_otel_span, proxy_logging_obj response: Final = await _dictable_table(TeamMembershipRepository(prisma_client)).find_unique( where={"user_id_team_id": {"user_id": user_id, "team_id": team_id}}, From ff023906855c96cb1d2a65e4fc02cd9882302761 Mon Sep 17 00:00:00 2001 From: yassin Date: Mon, 14 Sep 2026 19:26:14 +0000 Subject: [PATCH 09/13] test(auth): drop docstrings from team membership tests Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/test_litellm/proxy/auth/test_auth_checks.py | 6 ------ 1 file changed, 6 deletions(-) diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index 04e41608b26..b843916debf 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -6456,7 +6456,6 @@ async def test_get_team_membership_reads_sentinel_as_no_membership_not_a_model() @pytest.mark.asyncio async def test_get_team_membership_coalesces_parallel_db_fetches(): - """Concurrent misses for the same member must share one Prisma round-trip.""" from litellm.proxy.auth.auth_checks import get_team_membership started = asyncio.Event() @@ -6496,7 +6495,6 @@ async def test_get_team_membership_coalesces_parallel_db_fetches(): @pytest.mark.asyncio async def test_common_checks_calls_get_team_membership_once_per_request(): - """Model-access, attribution, and member-budget must reuse one membership load.""" from fastapi import Request from litellm.proxy.auth.auth_checks import common_checks @@ -6543,7 +6541,6 @@ async def test_common_checks_calls_get_team_membership_once_per_request(): @pytest.mark.asyncio async def test_common_checks_skips_membership_load_when_no_check_reads_it(): - """A management route has no model and no budget gate, so the membership row is never loaded.""" from fastapi import Request from litellm.proxy.auth.auth_checks import common_checks @@ -6583,7 +6580,6 @@ async def test_common_checks_skips_membership_load_when_no_check_reads_it(): @pytest.mark.asyncio async def test_get_team_membership_db_error_returns_none_and_retries_next_call(): - """A Prisma failure reads as no membership, caches nothing, and the next call hits the DB again.""" from litellm.proxy.auth.auth_checks import get_team_membership from litellm.proxy.common_utils.user_api_key_cache import team_membership_reservation_cache_key @@ -6620,7 +6616,6 @@ async def test_get_team_membership_db_error_returns_none_and_retries_next_call() @pytest.mark.asyncio async def test_get_team_membership_string_prisma_client_returns_none(): - """Unit tests stub prisma_client as a string; the lookup fails and reads as no membership.""" from litellm.proxy.auth.auth_checks import get_team_membership result = await get_team_membership( @@ -6634,7 +6629,6 @@ async def test_get_team_membership_string_prisma_client_returns_none(): @pytest.mark.asyncio async def test_get_team_membership_waiter_cancel_does_not_cancel_shared_load(): - """Cancelling one coalesced waiter must not cancel the shared Prisma load.""" from litellm.proxy.auth.auth_checks import get_team_membership started = asyncio.Event() From 4ac168b3fcb440a7e0af0e3b919ad6880baf2fd6 Mon Sep 17 00:00:00 2001 From: yassin Date: Mon, 14 Sep 2026 20:13:39 +0000 Subject: [PATCH 10/13] fix(auth): drop in-flight membership load on invalidation so it cannot repopulate the cache Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/auth/auth_checks.py | 8 ++- .../proxy/auth/test_auth_checks.py | 50 +++++++++++++++++++ 2 files changed, 56 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index fa85400f4e0..43d840130b6 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -2202,8 +2202,11 @@ async def _fetch_team_membership_from_db( where={"user_id_team_id": {"user_id": user_id, "team_id": team_id}}, include={"litellm_budget_table": True}, ) + membership: Final = None if response is None else LiteLLM_TeamMembership.model_validate(response.dict()) _key: Final = team_membership_reservation_cache_key(user_id=user_id, team_id=team_id) - if response is None: + if _team_membership_inflight.get(_key) is not asyncio.current_task(): + return membership + if membership is None: await user_api_key_cache.async_set_cache( key=_key, value=NO_TEAM_MEMBERSHIP_SENTINEL, @@ -2211,7 +2214,6 @@ async def _fetch_team_membership_from_db( ) return None - membership: Final = LiteLLM_TeamMembership.model_validate(response.dict()) await user_api_key_cache.async_set_cache( key=_key, value=membership, @@ -2762,6 +2764,8 @@ async def invalidate_team_member_spend_state( publish_auth_cache_invalidation, ) + _team_membership_inflight.pop(team_membership_reservation_cache_key(user_id=user_id, team_id=team_id), None) + if new_spend is not None: from litellm.proxy.proxy_server import SPEND_DB_FLOOR_CACHE_TTL_SECONDS, spend_counter_cache diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index b843916debf..986d53b80b0 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -6493,6 +6493,56 @@ async def test_get_team_membership_coalesces_parallel_db_fetches(): mock_prisma_client.db.litellm_teammembership.find_unique.assert_awaited_once() +@pytest.mark.asyncio +async def test_get_team_membership_invalidation_mid_flight_discards_stale_load(): + from litellm.proxy._types import LiteLLM_TeamMembership + from litellm.proxy.auth.auth_checks import get_team_membership, invalidate_team_member_spend_state + from litellm.proxy.common_utils.cache_pydantic_utils import CacheCodec + from litellm.proxy.common_utils.user_api_key_cache import team_membership_reservation_cache_key + + started = asyncio.Event() + release_stale = asyncio.Event() + release_fresh = asyncio.Event() + loads = iter((("budget-old", release_stale), ("budget-new", release_fresh))) + + async def _find_unique(*args, **kwargs): + budget_id, release = next(loads) + row = MagicMock() + row.dict = lambda: {"user_id": "u-inv", "team_id": "t-inv", "spend": 1.0, "budget_id": budget_id} + started.set() + await release.wait() + return row + + mock_prisma_client = MagicMock() + mock_prisma_client.db.litellm_teammembership.find_unique = AsyncMock(side_effect=_find_unique) + cache = UserApiKeyCache() + + async def _load(): + return await get_team_membership( + user_id="u-inv", team_id="t-inv", prisma_client=mock_prisma_client, user_api_key_cache=cache + ) + + stale = asyncio.create_task(_load()) + await started.wait() + await invalidate_team_member_spend_state(user_id="u-inv", team_id="t-inv", user_api_key_cache=cache) + started.clear() + fresh = asyncio.create_task(_load()) + await asyncio.wait_for(started.wait(), timeout=2) + release_fresh.set() + fresh_result = await fresh + release_stale.set() + stale_result = await stale + + assert stale_result is not None and stale_result.budget_id == "budget-old" + assert fresh_result is not None and fresh_result.budget_id == "budget-new" + assert mock_prisma_client.db.litellm_teammembership.find_unique.await_count == 2 + cached = CacheCodec.deserialize( + await cache.async_get_cache(key=team_membership_reservation_cache_key(user_id="u-inv", team_id="t-inv")), + model_type=LiteLLM_TeamMembership, + ) + assert cached is not None and cached.budget_id == "budget-new" + + @pytest.mark.asyncio async def test_common_checks_calls_get_team_membership_once_per_request(): from fastapi import Request From 91c964a338c759cb5ea69403d90a12696795d816 Mon Sep 17 00:00:00 2001 From: yassin Date: Mon, 14 Sep 2026 20:46:09 +0000 Subject: [PATCH 11/13] fix(auth): evict the membership cache entry when invalidation lands during the cache write Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/auth/auth_checks.py | 15 +++++---- .../proxy/auth/test_auth_checks.py | 32 +++++++++++++++++++ 2 files changed, 40 insertions(+), 7 deletions(-) diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 43d840130b6..22644bcd207 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -2212,13 +2212,14 @@ async def _fetch_team_membership_from_db( value=NO_TEAM_MEMBERSHIP_SENTINEL, ttl=get_management_object_ttl(user_api_key_cache), ) - return None - - await user_api_key_cache.async_set_cache( - key=_key, - value=membership, - model_type=LiteLLM_TeamMembership, - ) + else: + await user_api_key_cache.async_set_cache( + key=_key, + value=membership, + model_type=LiteLLM_TeamMembership, + ) + if _team_membership_inflight.get(_key) is not asyncio.current_task(): + await user_api_key_cache.async_delete_cache(key=_key) return membership diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index 986d53b80b0..61f0675aa1e 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -6543,6 +6543,38 @@ async def test_get_team_membership_invalidation_mid_flight_discards_stale_load() assert cached is not None and cached.budget_id == "budget-new" +@pytest.mark.asyncio +async def test_get_team_membership_invalidation_during_cache_write_evicts_stale_entry(): + from litellm.proxy.auth.auth_checks import get_team_membership, invalidate_team_member_spend_state + from litellm.proxy.common_utils.user_api_key_cache import team_membership_reservation_cache_key + + write_started = asyncio.Event() + release_write = asyncio.Event() + + class _SlowWriteCache(UserApiKeyCache): + async def async_set_cache(self, key, value, local_only=False, **kwargs): + write_started.set() + await release_write.wait() + return await super().async_set_cache(key, value, local_only=local_only, **kwargs) + + row = MagicMock() + row.dict = lambda: {"user_id": "u-w", "team_id": "t-w", "spend": 1.0, "budget_id": "budget-old"} + mock_prisma_client = MagicMock() + mock_prisma_client.db.litellm_teammembership.find_unique = AsyncMock(return_value=row) + cache = _SlowWriteCache() + + stale = asyncio.create_task( + get_team_membership(user_id="u-w", team_id="t-w", prisma_client=mock_prisma_client, user_api_key_cache=cache) + ) + await asyncio.wait_for(write_started.wait(), timeout=2) + await invalidate_team_member_spend_state(user_id="u-w", team_id="t-w", user_api_key_cache=cache) + release_write.set() + stale_result = await stale + + assert stale_result is not None and stale_result.budget_id == "budget-old" + assert await cache.async_get_cache(key=team_membership_reservation_cache_key(user_id="u-w", team_id="t-w")) is None + + @pytest.mark.asyncio async def test_common_checks_calls_get_team_membership_once_per_request(): from fastapi import Request From ffeea30f23d4de5dbd7554fa3a0518ab727e168a Mon Sep 17 00:00:00 2001 From: yassin Date: Mon, 14 Sep 2026 20:55:22 +0000 Subject: [PATCH 12/13] test(auth): cover a stale membership write landing after a fresh reload Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../proxy/auth/test_auth_checks.py | 45 +++++++++++++++++++ 1 file changed, 45 insertions(+) diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index 61f0675aa1e..96e9958a55d 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -6575,6 +6575,51 @@ async def test_get_team_membership_invalidation_during_cache_write_evicts_stale_ assert await cache.async_get_cache(key=team_membership_reservation_cache_key(user_id="u-w", team_id="t-w")) is None +@pytest.mark.asyncio +async def test_get_team_membership_stale_write_finishing_after_fresh_load_never_serves_old_row(): + from litellm.proxy.auth.auth_checks import get_team_membership, invalidate_team_member_spend_state + + write_started = asyncio.Event() + release_stale_write = asyncio.Event() + stale_writes = iter((release_stale_write,)) + + class _SlowFirstWriteCache(UserApiKeyCache): + async def async_set_cache(self, key, value, local_only=False, **kwargs): + release = next(stale_writes, None) + if release is not None: + write_started.set() + await release.wait() + return await super().async_set_cache(key, value, local_only=local_only, **kwargs) + + rows = iter(("budget-old", "budget-new", "budget-new")) + + async def _find_unique(*args, **kwargs): + row = MagicMock() + row.dict = lambda: {"user_id": "u-sw", "team_id": "t-sw", "spend": 1.0, "budget_id": next(rows)} + return row + + mock_prisma_client = MagicMock() + mock_prisma_client.db.litellm_teammembership.find_unique = AsyncMock(side_effect=_find_unique) + cache = _SlowFirstWriteCache() + + async def _load(): + return await get_team_membership( + user_id="u-sw", team_id="t-sw", prisma_client=mock_prisma_client, user_api_key_cache=cache + ) + + stale = asyncio.create_task(_load()) + await asyncio.wait_for(write_started.wait(), timeout=2) + await invalidate_team_member_spend_state(user_id="u-sw", team_id="t-sw", user_api_key_cache=cache) + fresh_result = await _load() + release_stale_write.set() + stale_result = await stale + after_result = await _load() + + assert fresh_result is not None and fresh_result.budget_id == "budget-new" + assert stale_result is not None and stale_result.budget_id == "budget-old" + assert after_result is not None and after_result.budget_id == "budget-new" + + @pytest.mark.asyncio async def test_common_checks_calls_get_team_membership_once_per_request(): from fastapi import Request From db8dfe93a5a74fad5c4ba3f9fabaa5de39ad09ca Mon Sep 17 00:00:00 2001 From: yassin Date: Mon, 14 Sep 2026 21:10:43 +0000 Subject: [PATCH 13/13] fix(auth): wait for the in-flight membership load before evicting its cache key on invalidation Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/auth/auth_checks.py | 10 +- .../proxy/auth/test_auth_checks.py | 95 ++++++------------- 2 files changed, 36 insertions(+), 69 deletions(-) diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 22644bcd207..13605d7dc7b 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -2204,8 +2204,6 @@ async def _fetch_team_membership_from_db( ) membership: Final = None if response is None else LiteLLM_TeamMembership.model_validate(response.dict()) _key: Final = team_membership_reservation_cache_key(user_id=user_id, team_id=team_id) - if _team_membership_inflight.get(_key) is not asyncio.current_task(): - return membership if membership is None: await user_api_key_cache.async_set_cache( key=_key, @@ -2218,8 +2216,6 @@ async def _fetch_team_membership_from_db( value=membership, model_type=LiteLLM_TeamMembership, ) - if _team_membership_inflight.get(_key) is not asyncio.current_task(): - await user_api_key_cache.async_delete_cache(key=_key) return membership @@ -2765,7 +2761,11 @@ async def invalidate_team_member_spend_state( publish_auth_cache_invalidation, ) - _team_membership_inflight.pop(team_membership_reservation_cache_key(user_id=user_id, team_id=team_id), None) + inflight: Final[object] = _team_membership_inflight.pop( + team_membership_reservation_cache_key(user_id=user_id, team_id=team_id), None + ) + if isinstance(inflight, asyncio.Task) and inflight is not asyncio.current_task(): + await asyncio.wait((inflight,)) if new_spend is not None: from litellm.proxy.proxy_server import SPEND_DB_FLOOR_CACHE_TTL_SECONDS, spend_counter_cache diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index 96e9958a55d..5f87f2def93 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -6494,7 +6494,7 @@ async def test_get_team_membership_coalesces_parallel_db_fetches(): @pytest.mark.asyncio -async def test_get_team_membership_invalidation_mid_flight_discards_stale_load(): +async def test_get_team_membership_invalidation_waits_for_in_flight_load_then_evicts_it(): from litellm.proxy._types import LiteLLM_TeamMembership from litellm.proxy.auth.auth_checks import get_team_membership, invalidate_team_member_spend_state from litellm.proxy.common_utils.cache_pydantic_utils import CacheCodec @@ -6502,20 +6502,21 @@ async def test_get_team_membership_invalidation_mid_flight_discards_stale_load() started = asyncio.Event() release_stale = asyncio.Event() - release_fresh = asyncio.Event() - loads = iter((("budget-old", release_stale), ("budget-new", release_fresh))) + rows = iter(("budget-old", "budget-new")) async def _find_unique(*args, **kwargs): - budget_id, release = next(loads) + budget_id = next(rows) row = MagicMock() row.dict = lambda: {"user_id": "u-inv", "team_id": "t-inv", "spend": 1.0, "budget_id": budget_id} - started.set() - await release.wait() + if budget_id == "budget-old": + started.set() + await release_stale.wait() return row mock_prisma_client = MagicMock() mock_prisma_client.db.litellm_teammembership.find_unique = AsyncMock(side_effect=_find_unique) cache = UserApiKeyCache() + _key = team_membership_reservation_cache_key(user_id="u-inv", team_id="t-inv") async def _load(): return await get_team_membership( @@ -6523,24 +6524,28 @@ async def test_get_team_membership_invalidation_mid_flight_discards_stale_load() ) stale = asyncio.create_task(_load()) - await started.wait() - await invalidate_team_member_spend_state(user_id="u-inv", team_id="t-inv", user_api_key_cache=cache) - started.clear() - fresh = asyncio.create_task(_load()) await asyncio.wait_for(started.wait(), timeout=2) - release_fresh.set() - fresh_result = await fresh - release_stale.set() - stale_result = await stale + invalidation = asyncio.create_task( + invalidate_team_member_spend_state(user_id="u-inv", team_id="t-inv", user_api_key_cache=cache) + ) + for _ in range(5): + await asyncio.sleep(0) + assert not invalidation.done() + release_stale.set() + await asyncio.wait_for(invalidation, timeout=2) + stale_result = await stale assert stale_result is not None and stale_result.budget_id == "budget-old" + assert await cache.async_get_cache(key=_key) is None + + fresh_result = await _load() assert fresh_result is not None and fresh_result.budget_id == "budget-new" assert mock_prisma_client.db.litellm_teammembership.find_unique.await_count == 2 - cached = CacheCodec.deserialize( - await cache.async_get_cache(key=team_membership_reservation_cache_key(user_id="u-inv", team_id="t-inv")), - model_type=LiteLLM_TeamMembership, - ) + cached = CacheCodec.deserialize(await cache.async_get_cache(key=_key), model_type=LiteLLM_TeamMembership) assert cached is not None and cached.budget_id == "budget-new" + again = await _load() + assert again is not None and again.budget_id == "budget-new" + assert mock_prisma_client.db.litellm_teammembership.find_unique.await_count == 2 @pytest.mark.asyncio @@ -6567,59 +6572,21 @@ async def test_get_team_membership_invalidation_during_cache_write_evicts_stale_ get_team_membership(user_id="u-w", team_id="t-w", prisma_client=mock_prisma_client, user_api_key_cache=cache) ) await asyncio.wait_for(write_started.wait(), timeout=2) - await invalidate_team_member_spend_state(user_id="u-w", team_id="t-w", user_api_key_cache=cache) + invalidation = asyncio.create_task( + invalidate_team_member_spend_state(user_id="u-w", team_id="t-w", user_api_key_cache=cache) + ) + for _ in range(5): + await asyncio.sleep(0) + assert not invalidation.done() + release_write.set() + await asyncio.wait_for(invalidation, timeout=2) stale_result = await stale assert stale_result is not None and stale_result.budget_id == "budget-old" assert await cache.async_get_cache(key=team_membership_reservation_cache_key(user_id="u-w", team_id="t-w")) is None -@pytest.mark.asyncio -async def test_get_team_membership_stale_write_finishing_after_fresh_load_never_serves_old_row(): - from litellm.proxy.auth.auth_checks import get_team_membership, invalidate_team_member_spend_state - - write_started = asyncio.Event() - release_stale_write = asyncio.Event() - stale_writes = iter((release_stale_write,)) - - class _SlowFirstWriteCache(UserApiKeyCache): - async def async_set_cache(self, key, value, local_only=False, **kwargs): - release = next(stale_writes, None) - if release is not None: - write_started.set() - await release.wait() - return await super().async_set_cache(key, value, local_only=local_only, **kwargs) - - rows = iter(("budget-old", "budget-new", "budget-new")) - - async def _find_unique(*args, **kwargs): - row = MagicMock() - row.dict = lambda: {"user_id": "u-sw", "team_id": "t-sw", "spend": 1.0, "budget_id": next(rows)} - return row - - mock_prisma_client = MagicMock() - mock_prisma_client.db.litellm_teammembership.find_unique = AsyncMock(side_effect=_find_unique) - cache = _SlowFirstWriteCache() - - async def _load(): - return await get_team_membership( - user_id="u-sw", team_id="t-sw", prisma_client=mock_prisma_client, user_api_key_cache=cache - ) - - stale = asyncio.create_task(_load()) - await asyncio.wait_for(write_started.wait(), timeout=2) - await invalidate_team_member_spend_state(user_id="u-sw", team_id="t-sw", user_api_key_cache=cache) - fresh_result = await _load() - release_stale_write.set() - stale_result = await stale - after_result = await _load() - - assert fresh_result is not None and fresh_result.budget_id == "budget-new" - assert stale_result is not None and stale_result.budget_id == "budget-old" - assert after_result is not None and after_result.budget_id == "budget-new" - - @pytest.mark.asyncio async def test_common_checks_calls_get_team_membership_once_per_request(): from fastapi import Request