From 56ed05418eb6f7bc1d32564cb969924fe0be0edd Mon Sep 17 00:00:00 2001
From: Ishaan Jaff <ishaan@berri.ai>
Date: Wed, 24 Jun 2026 13:27:13 -0700
Subject: [PATCH] perf(ai-gateway): key cache on by default (TTL 60s); set
 LITELLM_AUTH_CACHE_TTL_SECS=0 for strict per-connection verification

---
 .../crates/ai-gateway/src/auth/cache.rs       | 40 +++++++++++++++----
 1 file changed, 32 insertions(+), 8 deletions(-)

diff --git a/litellm-rust/crates/ai-gateway/src/auth/cache.rs b/litellm-rust/crates/ai-gateway/src/auth/cache.rs
index a4b7f9b5bc8..bca3b7143d8 100644
--- a/litellm-rust/crates/ai-gateway/src/auth/cache.rs
+++ b/litellm-rust/crates/ai-gateway/src/auth/cache.rs
@@ -1,11 +1,17 @@
-//! A bounded, in-memory cache of verified keys.
+//! A bounded, in-memory cache of verified keys — **on by default, for performance.**
 //!
-//! This bounds two things at once: **memory** (at most [`MAX_ENTRIES`] entries,
-//! FIFO-evicted) and **staleness** (each entry expires after [`KeyCache::ttl`], so
-//! a revoked or re-budgeted key is re-verified within the TTL window). Both are
-//! deliberate: the gateway trades a small staleness window for not hitting the
-//! auth backend on every request.
+//! Auth verification is a control-plane round-trip; caching it keeps the gateway
+//! from hammering the proxy's verify endpoint on every request under load — the
+//! same reason the LiteLLM proxy caches key-auth in-process. The TTL is the
+//! perf/freshness knob: a revoked, blocked, or re-budgeted key is re-verified
+//! within the TTL window (a deliberate, bounded staleness, identical in spirit to
+//! the proxy's own ~60s auth cache).
 //!
+//! For **strict per-connection enforcement** on billable routes (zero staleness —
+//! every connection re-verifies, budget/rate-limit checked each time), set
+//! `LITELLM_AUTH_CACHE_TTL_SECS=0` to disable caching.
+//!
+//! Bounds: at most [`MAX_ENTRIES`] entries (FIFO-evicted) plus TTL expiry.
 //! Std-only (no external cache crate): a `Mutex<HashMap>` for lookup plus a
 //! `VecDeque` recording insertion order for FIFO eviction.
 
@@ -18,9 +24,11 @@ use crate::auth::UserApiKeyAuth;
 /// Hard cap on cached entries; the 201st insert evicts the oldest. Keeps memory
 /// bounded regardless of how many distinct keys hit the gateway.
 const MAX_ENTRIES: usize = 200;
-/// Default time-to-live when `LITELLM_AUTH_CACHE_TTL_SECS` is unset/invalid.
+/// Default TTL when `LITELLM_AUTH_CACHE_TTL_SECS` is unset/invalid. On by default
+/// for performance; matches the proxy's own auth-cache window. Set `0` to disable
+/// (strict per-connection verification).
 const DEFAULT_TTL_SECS: u64 = 60;
-/// Env var overriding the entry TTL (seconds).
+/// Env var overriding the entry TTL (seconds). Set `0` to disable caching.
 const TTL_ENV: &str = "LITELLM_AUTH_CACHE_TTL_SECS";
 
 struct Inner {
@@ -61,6 +69,10 @@ impl KeyCache {
     /// Return the cached identity for `hash`, or `None`. A present-but-expired
     /// entry is evicted and treated as a miss, so the caller re-verifies.
     pub fn get(&self, hash: &[u8; 32]) -> Option<UserApiKeyAuth> {
+        // TTL 0 ⇒ caching disabled: always a miss, so every connection re-verifies.
+        if self.ttl.is_zero() {
+            return None;
+        }
         let mut inner = self.inner.lock().expect("auth cache mutex poisoned");
         let expired = match inner.entries.get(hash) {
             Some((_, inserted)) => inserted.elapsed() >= self.ttl,
@@ -78,6 +90,10 @@ impl KeyCache {
 
     /// Insert (or refresh) an entry, evicting the oldest if at capacity.
     pub fn insert(&self, hash: [u8; 32], auth: UserApiKeyAuth) {
+        // TTL 0 ⇒ caching disabled: don't store anything.
+        if self.ttl.is_zero() {
+            return;
+        }
         let mut inner = self.inner.lock().expect("auth cache mutex poisoned");
 
         // Reconcile `order` so it tracks exactly the live entries: drop this key's
@@ -154,4 +170,12 @@ mod tests {
         std::thread::sleep(Duration::from_millis(25));
         assert!(cache.get(&hash(1)).is_none());
     }
+
+    #[test]
+    fn ttl_zero_disables_caching() {
+        // The default: every lookup is a miss, so callers always re-verify.
+        let cache = KeyCache::with_ttl(Duration::from_secs(0));
+        cache.insert(hash(1), UserApiKeyAuth::default());
+        assert!(cache.get(&hash(1)).is_none());
+    }
 }