refactor: split api key and tier identity caches

Author: suguanYang

apps/api/app/core/dependencies.py

@@ -1,23 +1,22 @@
-import hashlib
 import threading
 from datetime import timedelta
 from fnmatch import fnmatch
 from typing import Any
 
 import jwt
 from app.services.auth.api_key_service import APIKeyService
-from app.services.rate_limit.identity_cache import identity_cache
 from fastapi import Depends, Header, Request
 from jwt import PyJWKClient
 from loguru import logger
 from sqlalchemy.ext.asyncio import AsyncSession
 
-from shared.core.config import redis_pool_manager, settings
+from shared.core.config import settings
 from shared.core.database import get_db
 from shared.core.exceptions.domain_exceptions import (
     AuthException,
     PermissionDeniedException,
 )
+from shared.utils.api_keys import is_api_key_token
 
 # Standard JWKS endpoint path (fixed, following OpenID Connect convention)
 JWKS_ENDPOINT_PATH = "/api/auth/jwks"
@@ -184,39 +183,16 @@ async def get_current_user_id(
     route_path = _get_route_path(request)
 
     # Mode 1: API Key verification (for external clients)
-    if token.startswith("sk_"):
-        # Check identity cache first — skip DB on cache hit
-        api_key_hash = hashlib.sha256(token.encode()).hexdigest()
-        try:
-            cached = await identity_cache.get_cached_identity(
-                redis_pool_manager.get_redis_service(),
-                identity_cache._apikey_key(api_key_hash),
-            )
-            if cached is not None:
-                cached_user_id = cached.get("user_id")
-                cached_user_tier = cached.get("user_tier")
-                if cached_user_id and isinstance(cached_user_tier, str):
-                    request.state.cached_user_tier = cached_user_tier
-                    request.state.cached_identity_hit = True
-                    request.state.user_id = cached_user_id
-                    _enforce_guest_api_key_scope(route_path, cached_user_tier)
-                    return cached_user_id
-        except PermissionDeniedException:
-            raise
-        except Exception:
-            pass  # Fall through to DB validation
-
-        # Cache miss — validate via DB
+    if is_api_key_token(token):
         api_key_service = APIKeyService()
-        identity = await api_key_service.validate_api_key_identity(db, token)
+        identity = await api_key_service.get_identity(db, token)
         if identity:
             request.state.cached_user_tier = identity.user_tier
-            request.state.cached_identity_hit = False
             request.state.user_id = identity.user_id
             _enforce_guest_api_key_scope(route_path, identity.user_tier)
             return identity.user_id
-        else:
-            raise AuthException(user_message="Invalid API Key")
+
+        raise AuthException(user_message="Invalid API Key")
 
     # Mode 2: JWT verification (for Dashboard/Internal)
     return decode_jwt_token(token)

apps/api/app/services/auth/api_key_identity_cache.py

@@ -0,0 +1,124 @@
+"""Redis-backed API-key authentication user cache."""
+
+from __future__ import annotations
+
+import json
+from typing import TYPE_CHECKING
+
+from loguru import logger
+
+if TYPE_CHECKING:
+    from shared.services.redis.redis_service import RedisService
+
+_API_KEY_MAX_TTL_SECONDS: int = 3600
+
+
+class APIKeyIdentityCache:
+    """Cache validated API-key user IDs by API-key lookup hash."""
+
+    @staticmethod
+    def get_cache_key(api_key_hash: str) -> str:
+        """Return the Redis key for an API-key hash."""
+        return f"identity:apikey:{api_key_hash}"
+
+    @staticmethod
+    def get_reverse_key(user_id: str) -> str:
+        """Return the reverse-index Redis key for a user."""
+        return f"identity:apikeys:{user_id}"
+
+    async def get_user_id(
+        self,
+        redis: RedisService,
+        api_key_hash: str,
+    ) -> str | None:
+        """Return cached user_id for an API key."""
+        try:
+            raw_user_id: object = await redis.get(self.get_cache_key(api_key_hash))
+            return self._coerce_user_id(raw_user_id)
+        except Exception:
+            logger.warning("api_key_identity_cache: failed to read user")
+            return None
+
+    async def set_user_id(
+        self,
+        redis: RedisService,
+        api_key_hash: str,
+        user_id: str,
+        ttl_seconds: int,
+    ) -> None:
+        """Cache a validated API-key user ID."""
+        effective_ttl_seconds: int = min(_API_KEY_MAX_TTL_SECONDS, ttl_seconds)
+        cache_key: str = self.get_cache_key(api_key_hash)
+        reverse_key: str = self.get_reverse_key(user_id)
+
+        try:
+            await redis.set(cache_key, user_id, ttl=effective_ttl_seconds)
+            await redis.sadd(reverse_key, api_key_hash)
+            current_ttl_seconds: int = await redis.ttl(reverse_key)
+            if (
+                current_ttl_seconds in (-2, -1)
+                or current_ttl_seconds < effective_ttl_seconds
+            ):
+                await redis.expire(reverse_key, effective_ttl_seconds)
+        except Exception:
+            logger.warning(
+                "api_key_identity_cache: failed to set user for user_id={}",
+                user_id,
+            )
+
+    async def invalidate_api_key(
+        self,
+        redis: RedisService,
+        user_id: str,
+        api_key_hash: str,
+    ) -> None:
+        """Delete one API-key identity cache entry."""
+        try:
+            await redis.delete(self.get_cache_key(api_key_hash))
+            await redis.srem(self.get_reverse_key(user_id), api_key_hash)
+        except Exception:
+            logger.warning(
+                "api_key_identity_cache: failed to invalidate identity for user_id={}",
+                user_id,
+            )
+
+    async def invalidate_user(
+        self,
+        redis: RedisService,
+        user_id: str,
+    ) -> None:
+        """Delete all API-key identity cache entries for a user."""
+        try:
+            reverse_key: str = self.get_reverse_key(user_id)
+            api_key_hashes: set[object] = await redis.smembers(reverse_key)
+            for api_key_hash in api_key_hashes:
+                await redis.delete(self.get_cache_key(str(api_key_hash)))
+            await redis.delete(reverse_key)
+        except Exception:
+            logger.warning(
+                "api_key_identity_cache: failed to invalidate user_id={}",
+                user_id,
+            )
+
+    def _coerce_user_id(self, raw_user_id: object) -> str | None:
+        """Return a typed user ID from current or legacy Redis values."""
+        if isinstance(raw_user_id, str):
+            try:
+                parsed_user_id: object = json.loads(raw_user_id)
+            except json.JSONDecodeError:
+                return raw_user_id
+        else:
+            parsed_user_id = raw_user_id
+
+        if isinstance(parsed_user_id, str):
+            return parsed_user_id
+
+        if isinstance(parsed_user_id, dict):
+            legacy_user_id: object = parsed_user_id.get("user_id")
+            if isinstance(legacy_user_id, str):
+                return legacy_user_id
+
+        return None
+
+
+api_key_identity_cache = APIKeyIdentityCache()