Coverage for .venv/lib/python3.13/site-packages/litellm/proxy/common_utils/user_api_key_cache.py: 64%

1from __future__ import annotations 

2 

3import asyncio 

4import re 

5from collections.abc import Sequence 

6from typing import TYPE_CHECKING, Any, Final, TypeVar, cast, overload 

7 

8from pydantic import BaseModel 

9 

10from litellm._logging import verbose_proxy_logger 

11from litellm.caching.dual_cache import DualCache 

12from litellm.caching.in_memory_cache import InMemoryCache 

13from litellm.caching.redis_cache import RedisCache 

14from litellm.constants import DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL 

15from litellm.proxy.common_utils.cache_pydantic_utils import CacheCodec 

16 

17if TYPE_CHECKING: 17 ↛ 18line 17 didn't jump to line 18 because the condition on line 17 was never true

18 from opentelemetry.trace import Span 

19 

20T = TypeVar("T", bound=BaseModel) 

21 

22_HASHED_TOKEN_CACHE_KEY: Final = re.compile(r"[0-9a-f]{64}") 

23 

24 

25def is_user_key_cache_key(key: str) -> bool: 

26 """Only user-key objects are cached under a bare ``hash_token`` digest; every other object uses a prefixed key.""" 

27 return _HASHED_TOKEN_CACHE_KEY.fullmatch(key) is not None 

28 

29 

30class UserApiKeyCache(DualCache): 

31 """ 

32 DualCache wrapper for UserAPIKeyAuth-like payloads. 

33 

34 Stores a Redis-safe JSON payload in BOTH in-memory and Redis to avoid 

35 "memory returns BaseModel, Redis returns dict" format drift. 

36 

37 When ``model_type`` is provided: 

38 - writes are serialized via ``CacheCodec.serialize(..., model_type=...)`` 

39 - reads are deserialized via ``CacheCodec.deserialize(..., model_type)`` 

40 and return ``Optional[T]``: the model on success, ``None`` on cache miss 

41 **or** if the cached payload fails validation (schema drift). On 

42 validation failure after a cache hit, an error line is emitted via 

43 ``verbose_proxy_logger``. 

44 

45 When ``model_type`` is omitted, the interface behaves like ``DualCache``: 

46 raw cached payload is returned (dict/str/etc.). 

47 

48 ``async_set_cache_pipeline`` applies the same untyped Codec pass as omitting 

49 ``model_type`` on ``async_set_cache`` (so ``BaseModel`` rows are dumped before Redis). 

50 

51 User-key objects (see ``is_user_key_cache_key``) live in their own in-memory partition, 

52 ``key_object_cache``, so churn in the other management objects cannot evict them. Both 

53 partitions share the same Redis backend and TTL settings. 

54 

55 ``get_cache`` / ``async_get_cache`` overloads and implementations must be contiguous 

56 (no other methods in between) so mypy resolves ``@overload`` + implementation correctly. 

57 """ 

58 

59 def __init__( 

60 self, 

61 in_memory_cache: InMemoryCache | None = None, 

62 redis_cache: RedisCache | None = None, 

63 default_in_memory_ttl: float | None = None, 

64 default_redis_ttl: float | None = None, 

65 key_object_in_memory_cache: InMemoryCache | None = None, 

66 ) -> None: 

67 super().__init__( 

68 in_memory_cache=in_memory_cache, 

69 redis_cache=redis_cache, 

70 default_in_memory_ttl=default_in_memory_ttl, 

71 default_redis_ttl=default_redis_ttl, 

72 ) 

73 self.key_object_cache: Final = DualCache( 

74 in_memory_cache=key_object_in_memory_cache or InMemoryCache(), 

75 redis_cache=redis_cache, 

76 default_in_memory_ttl=default_in_memory_ttl, 

77 default_redis_ttl=default_redis_ttl, 

78 ) 

79 

80 def in_memory_cache_for(self, key: str) -> InMemoryCache: 

81 return self.key_object_cache.in_memory_cache if is_user_key_cache_key(key) else self.in_memory_cache 

82 

83 def update_cache_ttl(self, default_in_memory_ttl: float | None, default_redis_ttl: float | None) -> None: 

84 super().update_cache_ttl(default_in_memory_ttl=default_in_memory_ttl, default_redis_ttl=default_redis_ttl) 

85 self.key_object_cache.update_cache_ttl( 

86 default_in_memory_ttl=default_in_memory_ttl, default_redis_ttl=default_redis_ttl 

87 ) 

88 

89 def update_in_memory_max_size(self, max_size: int | None) -> None: 

90 super().update_in_memory_max_size(max_size) 

91 self.key_object_cache.update_in_memory_max_size(max_size) 

92 

93 def attach_redis_cache( 

94 self, redis_cache: RedisCache | None = None, *, default_redis_ttl: float | None = None 

95 ) -> None: 

96 super().attach_redis_cache(redis_cache, default_redis_ttl=default_redis_ttl) 

97 self.key_object_cache.attach_redis_cache(redis_cache, default_redis_ttl=default_redis_ttl) 

98 

99 @overload 

100 def get_cache( 100 ↛ exitline 100 didn't return from function 'get_cache' because

101 self, 

102 key: str, 

103 parent_otel_span: Span | None = None, 

104 local_only: bool = False, 

105 *, 

106 model_type: type[T], 

107 **kwargs: object, 

108 ) -> T | None: ... 

109 

110 @overload 

111 def get_cache( 111 ↛ exitline 111 didn't return from function 'get_cache' because

112 self, 

113 key: str, 

114 parent_otel_span: Span | None = None, 

115 local_only: bool = False, 

116 model_type: None = None, 

117 **kwargs: object, 

118 ) -> Any: ... 

119 

120 def get_cache( 

121 self, 

122 key: str, 

123 parent_otel_span: Span | None = None, 

124 local_only: bool = False, 

125 model_type: type[BaseModel] | None = None, 

126 **kwargs: object, 

127 ) -> object: 

128 if model_type is None and "model_type" in kwargs: 

129 model_type = cast(type[BaseModel] | None, kwargs.pop("model_type", None)) 

130 cached: Final = ( 

131 self.key_object_cache.get_cache(key=key, parent_otel_span=parent_otel_span, local_only=local_only, **kwargs) 

132 if is_user_key_cache_key(key) 

133 else super().get_cache(key=key, parent_otel_span=parent_otel_span, local_only=local_only, **kwargs) 

134 ) 

135 if model_type is None: 

136 return cached 

137 if cached is None: 

138 return None 

139 decoded: Final = CacheCodec.deserialize(cached, model_type=model_type) 

140 if decoded is None: 

141 verbose_proxy_logger.error( 

142 "UserApiKeyCache.get_cache failed to deserialize cached value for key=%r model_type=%s", 

143 key, 

144 getattr(model_type, "__name__", str(model_type)), 

145 ) 

146 return None 

147 return decoded 

148 

149 @overload 

150 async def async_get_cache( 150 ↛ exitline 150 didn't return from function 'async_get_cache' because

151 self, 

152 key: str, 

153 parent_otel_span: Span | None = None, 

154 local_only: bool = False, 

155 *, 

156 model_type: type[T], 

157 **kwargs: object, 

158 ) -> T | None: ... 

159 

160 @overload 

161 async def async_get_cache( 161 ↛ exitline 161 didn't return from function 'async_get_cache' because

162 self, 

163 key: str, 

164 parent_otel_span: Span | None = None, 

165 local_only: bool = False, 

166 model_type: None = None, 

167 **kwargs: object, 

168 ) -> Any: ... 

169 

170 async def async_get_cache( 

171 self, 

172 key: str, 

173 parent_otel_span: Span | None = None, 

174 local_only: bool = False, 

175 model_type: type[BaseModel] | None = None, 

176 **kwargs: object, 

177 ) -> object: 

178 if model_type is None and "model_type" in kwargs: 178 ↛ 179line 178 didn't jump to line 179 because the condition on line 178 was never true

179 model_type = cast(type[BaseModel] | None, kwargs.pop("model_type", None)) 

180 cached: Final = ( 

181 await self.key_object_cache.async_get_cache( 

182 key=key, parent_otel_span=parent_otel_span, local_only=local_only, **kwargs 

183 ) 

184 if is_user_key_cache_key(key) 

185 else await super().async_get_cache( 

186 key=key, parent_otel_span=parent_otel_span, local_only=local_only, **kwargs 

187 ) 

188 ) 

189 if model_type is None: 

190 return cached 

191 if cached is None: 

192 return None 

193 decoded: Final = CacheCodec.deserialize(cached, model_type=model_type) 

194 if decoded is None: 194 ↛ 195line 194 didn't jump to line 195 because the condition on line 194 was never true

195 verbose_proxy_logger.error( 

196 "UserApiKeyCache.async_get_cache failed to deserialize cached value for key=%r model_type=%s", 

197 key, 

198 getattr(model_type, "__name__", str(model_type)), 

199 ) 

200 return None 

201 return decoded 

202 

203 def set_cache(self, key: str | None, value: object, local_only: bool = False, **kwargs: object): 

204 model_type: Final = cast(type[BaseModel] | None, kwargs.pop("model_type", None)) 

205 payload: Final[object] = CacheCodec.serialize(value, model_type=model_type) 

206 if key is not None and is_user_key_cache_key(key): 

207 return self.key_object_cache.set_cache(key=key, value=payload, local_only=local_only, **kwargs) 

208 return super().set_cache(key=key, value=payload, local_only=local_only, **kwargs) 

209 

210 async def async_set_cache(self, key: str | None, value: object, local_only: bool = False, **kwargs: object): 

211 model_type: Final = cast(type[BaseModel] | None, kwargs.pop("model_type", None)) 

212 payload: Final[object] = CacheCodec.serialize(value, model_type=model_type) 

213 if key is not None and is_user_key_cache_key(key): 

214 return await self.key_object_cache.async_set_cache(key=key, value=payload, local_only=local_only, **kwargs) 

215 return await super().async_set_cache(key=key, value=payload, local_only=local_only, **kwargs) 

216 

217 def delete_cache(self, key: str) -> None: 

218 if is_user_key_cache_key(key): 

219 self.key_object_cache.delete_cache(key) 

220 return 

221 super().delete_cache(key) 

222 

223 async def async_delete_cache(self, key: str) -> None: 

224 if is_user_key_cache_key(key): 224 ↛ 225line 224 didn't jump to line 225 because the condition on line 224 was never true

225 await self.key_object_cache.async_delete_cache(key) 

226 return 

227 await super().async_delete_cache(key) 

228 

229 async def async_delete_cache_keys(self, keys: Sequence[str]) -> None: 

230 """Batch twin of ``async_delete_cache``, partitioned like 

231 ``async_set_cache_pipeline``. 

232 

233 Both partitions are cleared even when one raises, because a caller 

234 batching these has already committed the rows they cache. 

235 """ 

236 key_object_keys: Final = tuple(key for key in keys if is_user_key_cache_key(key)) 

237 other_keys: Final = tuple(key for key in keys if not is_user_key_cache_key(key)) 

238 outcomes: Final = await asyncio.gather( 

239 self.key_object_cache.async_delete_cache_keys(key_object_keys), 

240 super().async_delete_cache_keys(other_keys), 

241 return_exceptions=True, 

242 ) 

243 failed: Final = tuple(outcome for outcome in outcomes if isinstance(outcome, BaseException)) 

244 if failed: 

245 raise failed[0] 

246 

247 def flush_cache(self) -> None: 

248 super().flush_cache() 

249 self.key_object_cache.in_memory_cache.flush_cache() 

250 

251 async def async_set_cache_pipeline( 

252 self, cache_list: Sequence[tuple[str, object]], local_only: bool = False, **kwargs: object 

253 ) -> None: 

254 """ 

255 Batch writes with the same Codec boundary as ``async_set_cache`` without 

256 ``model_type``: ``BaseModel`` values become JSON-safe dicts; dicts/scalars unchanged. 

257 """ 

258 normalized: Final = tuple((key, CacheCodec.serialize(value, model_type=None)) for key, value in cache_list) 

259 key_object_entries: Final = tuple(entry for entry in normalized if is_user_key_cache_key(entry[0])) 

260 other_entries: Final = tuple(entry for entry in normalized if not is_user_key_cache_key(entry[0])) 

261 if key_object_entries: 261 ↛ 262line 261 didn't jump to line 262 because the condition on line 261 was never true

262 await self.key_object_cache.async_set_cache_pipeline( 

263 cache_list=key_object_entries, local_only=local_only, **kwargs 

264 ) 

265 if other_entries: 265 ↛ exitline 265 didn't return from function 'async_set_cache_pipeline' because the condition on line 265 was always true

266 await super().async_set_cache_pipeline(cache_list=other_entries, local_only=local_only, **kwargs) 

267 

268 

269#: Value cached under ``user_object_permission_id_cache_key`` when the user links no permission row, 

270#: so a human without an entitlement costs no DB read per request. Lives beside the key builder 

271#: because it is part of the same cache protocol: a reader that knows the key must know this value. 

272USER_NO_MCP_PERMISSION_SENTINEL: Final = "__user_no_mcp_permission__" 

273 

274 

275def user_object_permission_id_cache_key(user_id: str) -> str: 

276 """Cache key for the ``user_id -> object_permission_id`` link. 

277 

278 Lives here rather than next to either user because two modules own the two halves: the MCP auth 

279 resolver writes it on read, and ``/user/update`` deletes it after changing the link. A key format 

280 duplicated across those two drifts silently, and the failure is an entitlement change that never 

281 takes effect. 

282 """ 

283 return f"user_object_permission_id:{user_id}" 

284 

285 

286def object_permission_cache_key(object_permission_id: str) -> str: 

287 """Cache key ``get_object_permission`` stores a permission row under.""" 

288 return f"object_permission_id:{object_permission_id}" 

289 

290 

291#: Cached under ``tag_registry_cache_key`` when the table exceeds ``TAG_REGISTRY_MAX_SIZE``: 

292#: registry unusable, fall back to the per-tag lookup. 

293TAG_REGISTRY_OVERFLOW_SENTINEL: Final = "__tag_registry_overflow__" 

294 

295 

296def tag_cache_key(tag_name: str) -> str: 

297 """Cache key one tag row is stored under; shared so its five reader/writer modules cannot drift.""" 

298 return f"tag:{tag_name}" 

299 

300 

301def tag_registry_cache_key() -> str: 

302 """Cache key for the set of tag names that exist in ``LiteLLM_TagTable``.""" 

303 return "tag_registry" 

304 

305 

306#: Cached under ``model_access_group_registry_cache_key`` when the table exceeds 

307#: ``MODEL_ACCESS_GROUP_REGISTRY_MAX_SIZE``: registry unusable, fall back to the per-group lookup. 

308MODEL_ACCESS_GROUP_REGISTRY_OVERFLOW_SENTINEL: Final = "__model_access_group_registry_overflow__" 

309 

310 

311def model_access_group_cache_key(access_group_name: str) -> str: 

312 """Cache key one model access group budget row is stored under; shared so auth, spend tracking and the management endpoints cannot drift.""" 

313 return f"model_access_group:{access_group_name}" 

314 

315 

316def model_access_group_registry_cache_key() -> str: 

317 """Cache key for the set of model access group names that have a budget row.""" 

318 return "model_access_group_registry" 

319 

320 

321def model_access_group_spend_counter_key(access_group_name: str) -> str: 

322 """Spend counter key for one model access group; shared so its four owners cannot drift. 

323 

324 The reservation path writes it up front, the cost callback writes it after the call, auth 

325 reads it to enforce ``max_budget``, and the reset job clears it on rollover. A copy that 

326 drifts in any one of them silently resets or reads a counter nobody else touches, which shows 

327 up as a budget that never trips or never resets. 

328 """ 

329 return f"spend:model_access_group:{access_group_name}" 

330 

331 

332def project_cache_key(project_id: str) -> str: 

333 return f"project_id:{project_id}" 

334 

335 

336def project_spend_counter_key(project_id: str) -> str: 

337 return f"spend:project:{project_id}" 

338 

339 

340#: Cached under ``end_user_restricted_registry_cache_key`` when the restricted set exceeds 

341#: ``END_USER_RESTRICTED_REGISTRY_MAX_SIZE``: registry unusable, fall back to the per-id fetch. 

342END_USER_RESTRICTED_REGISTRY_OVERFLOW_SENTINEL: Final = "__end_user_restricted_registry_overflow__" 

343 

344 

345def end_user_cache_key(end_user_id: str) -> str: 

346 """Cache key one end-user row is stored under; shared so auth and spend tracking cannot drift.""" 

347 return f"end_user_id:{end_user_id}" 

348 

349 

350def end_user_restricted_registry_cache_key() -> str: 

351 """Cache key for the set of end-user ids whose row carries a restriction auth enforces.""" 

352 return "end_user_restricted_registry" 

353 

354 

355def team_membership_auth_cache_key(team_id: str, user_id: str) -> str: 

356 """Cache key one team member's ``LiteLLM_TeamMembership`` row is stored under for the admission check.""" 

357 return f"{team_id}_{user_id}" 

358 

359 

360def team_membership_reservation_cache_key(user_id: str, team_id: str) -> str: 

361 """Cache key the pre-call budget reservation stores the same ``LiteLLM_TeamMembership`` row under. 

362 

363 Deliberately not unified with ``team_membership_auth_cache_key``: the two readers wrote independent 

364 keys before this file existed, so a fix that invalidates one must invalidate both explicitly rather 

365 than assume a single write is visible to both. 

366 """ 

367 return f"team_membership:{user_id}:{team_id}" 

368 

369 

370#: Cached under ``team_membership_reservation_cache_key`` when a member has no ``LiteLLM_TeamMembership`` 

371#: row, so a session-token member without a per-member budget costs no DB read per request. Lives beside 

372#: the key builder because it is part of the same cache protocol: every reader of the key must know that 

373#: a plain string here means "no row", distinct from a serialized membership. The two budget readers 

374#: already treat a non-model value as "no row", so they need no change to stay correct. 

375NO_TEAM_MEMBERSHIP_SENTINEL: Final = "__no_team_membership__" 

376 

377 

378def get_management_object_ttl(cache: DualCache) -> float: 

379 """ 

380 In-memory TTL for management-object cache writes (keys, teams, users, budgets, ...). 

381 

382 Honors ``general_settings.user_api_key_cache_ttl``, which ``proxy_server`` 

383 propagates onto ``default_in_memory_ttl`` at startup, and falls back to 

384 ``DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL`` when no default is configured. 

385 """ 

386 configured: Final[float | None] = getattr(cache, "default_in_memory_ttl", None) 

387 if configured is not None: 387 ↛ 389line 387 didn't jump to line 389 because the condition on line 387 was always true

388 return configured 

389 return DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL