fix: refactor

backend/auth/routers/oidc.py

@@ -13,7 +13,7 @@ from litellm.proxy.auth_v2.models import AuthMethod
 from ..services.redirects import safe_relay_state
 from litellm.proxy.auth_v2.resolvers import ProvisioningStore
 from litellm.proxy.auth_v2.security import AuthSecurity
-from litellm.proxy.auth_v2.sessions.schemas import OAuthTransaction, SessionState
+from litellm.proxy.auth_v2.sessions.types import OAuthTransaction, SessionState
 
 from ..services.oidc import mapped_claims, providers_by_key, user_from_userinfo
 from .dependencies import get_auth, get_oauth_registry

backend/auth/routers/saml.py

@@ -15,7 +15,7 @@ from litellm.proxy.auth_v2.authorization import filter_claim_roles
 from ..services.redirects import safe_relay_state
 from litellm.proxy.auth_v2.resolvers import ProvisioningStore
 from litellm.proxy.auth_v2.security import AuthSecurity
-from litellm.proxy.auth_v2.sessions.schemas import SessionState
+from litellm.proxy.auth_v2.sessions.types import SessionState
 
 from ..services.saml import (
     SAMLProtocolStore,

litellm/proxy/auth_v2/authenticators/session.py

@@ -11,12 +11,12 @@ from litellm.proxy.auth_v2.models import (
     CredentialRef,
     SecuritySchemeType,
 )
-from litellm.proxy.auth_v2.sessions import StateStore
-from litellm.proxy.auth_v2.sessions.schemas import SessionState
+from litellm.proxy.auth_v2.sessions import SessionStore
+from litellm.proxy.auth_v2.sessions.types import SessionState
 
 
 class SessionAuthenticator(Authenticator):
-    def __init__(self, cookie_name: str, store: "StateStore[SessionState]") -> None:
+    def __init__(self, cookie_name: str, store: "SessionStore[SessionState]") -> None:
         self._cookie_name = cookie_name
         self._store = store
 

litellm/proxy/auth_v2/network.py

@@ -34,13 +34,9 @@ def ip_in_trusted_proxies(ip: Optional[str], config: TrustedProxyConfig) -> bool
     return _ip_in_cidrs(ip, config.trusted_proxy_cidrs)
 
 
-def resolve_client_ip(
-    request: Request, config: TrustedProxyConfig
-) -> Tuple[Optional[str], bool]:
+def resolve_client_ip(request: Request, config: TrustedProxyConfig) -> Tuple[Optional[str], bool]:
     peer = request.client.host if request.client else None
-    if not config.use_forwarded_for or not _ip_in_cidrs(
-        peer, config.trusted_proxy_cidrs
-    ):
+    if not config.use_forwarded_for or not _ip_in_cidrs(peer, config.trusted_proxy_cidrs):
         return peer, False
     forwarded = request.headers.get("x-forwarded-for", "")
     hops = [h.strip() for h in forwarded.split(",") if h.strip()]
@@ -50,9 +46,7 @@ def resolve_client_ip(
     return peer, True
 
 
-def resolve_network_context(
-    request: Request, config: TrustedProxyConfig
-) -> NetworkContext:
+def resolve_network_context(request: Request, config: TrustedProxyConfig) -> NetworkContext:
     ip, via_proxy = resolve_client_ip(request, config)
     return NetworkContext(
         client_ip=ip,

litellm/proxy/auth_v2/security.py

@@ -1,8 +1,10 @@
+import os
 from typing import Annotated, Callable, List, Optional
 
 from fastapi import Request, Security
 from fastapi.security import SecurityScopes
 
+from litellm._redis import get_redis_async_client
 from litellm.proxy.auth_v2 import errors
 from litellm.proxy.auth_v2.authenticators import (
     Authenticator,
@@ -20,8 +22,34 @@ from litellm.proxy.auth_v2.authorization import (
 )
 from litellm.proxy.auth_v2.authenticators.session import SessionAuthenticator
 from litellm.proxy.auth_v2.resolvers import IdentityResolver
-from litellm.proxy.auth_v2.sessions import StateBackend, StateStore
-from litellm.proxy.auth_v2.sessions.schemas import OAuthTransaction, SessionState
+from litellm.proxy.auth_v2.sessions import (
+    InMemorySessionStore,
+    RedisSessionStore,
+    SessionStore,
+    SessionValue,
+)
+from litellm.proxy.auth_v2.sessions.types import OAuthTransaction, SessionState
+
+_REDIS_ENV_SIGNALS = (
+    "REDIS_URL",
+    "REDIS_HOST",
+    "REDIS_CLUSTER_NODES",
+    "REDIS_SENTINEL_NODES",
+)
+
+
+def _open_session_store(namespace: str, *, default_ttl: int) -> SessionStore[SessionValue]:
+    """Build the session/login-state store for ``namespace``.
+
+    Uses Redis when configured via the environment (required so state is shared
+    across pods in a multi-pod deployment); otherwise a process-local in-memory
+    store for single-process/dev. Chosen from the environment, not by probing
+    Redis, so a configured-but-unreachable Redis fails loudly on use rather than
+    silently stranding state on one pod.
+    """
+    if any(os.getenv(signal) for signal in _REDIS_ENV_SIGNALS):
+        return RedisSessionStore(get_redis_async_client(), namespace, default_ttl)
+    return InMemorySessionStore(namespace, default_ttl)
 
 
 def _combined_challenge(authenticators: List[Authenticator]) -> str:
@@ -56,16 +84,14 @@ class AuthSecurity:
         authorizer: Optional[Authorizer] = None,
         authenticators: Optional[List[Authenticator]] = None,
         basic_verifier: Optional[BasicAuthVerifier] = None,
-        state_backend: Optional[StateBackend] = None,
     ) -> None:
         self.config = config
         self.resolver = resolver
         self.authorizer = authorizer or RBACEngine(config.casbin_policy_path)
-        self._state = state_backend or StateBackend(None)
-        self.session_store: StateStore[SessionState] = self._state.store(
+        self.session_store: SessionStore[SessionState] = _open_session_store(
             "sessions", default_ttl=config.session.ttl_seconds
         )
-        self.oauth_txn_store: StateStore[OAuthTransaction] = self._state.store(
+        self.oauth_txn_store: SessionStore[OAuthTransaction] = _open_session_store(
             "oauth_txn", default_ttl=config.session.login_state_ttl
         )
         chain = (

litellm/proxy/auth_v2/sessions/__init__.py

@@ -1,15 +1,13 @@
-from litellm.proxy.auth_v2.sessions.factory import StateBackend
-from litellm.proxy.auth_v2.sessions.base import StateStore, StateValue
-from litellm.proxy.auth_v2.sessions.memory import InMemoryStateStore
-from litellm.proxy.auth_v2.sessions.redis import RedisStateStore
-from litellm.proxy.auth_v2.sessions.schemas import OAuthTransaction, SessionState
+from litellm.proxy.auth_v2.sessions.base import SessionStore, SessionValue
+from litellm.proxy.auth_v2.sessions.memory import InMemorySessionStore
+from litellm.proxy.auth_v2.sessions.redis import RedisSessionStore
+from litellm.proxy.auth_v2.sessions.types import OAuthTransaction, SessionState
 
 __all__ = [
-    "StateBackend",
-    "StateStore",
-    "StateValue",
-    "InMemoryStateStore",
-    "RedisStateStore",
+    "SessionStore",
+    "SessionValue",
+    "InMemorySessionStore",
+    "RedisSessionStore",
     "SessionState",
     "OAuthTransaction",
 ]

litellm/proxy/auth_v2/sessions/base.py

@@ -2,11 +2,11 @@ from __future__ import annotations
 
 from typing import Any, Mapping, Optional, Protocol, TypeVar, runtime_checkable
 
-StateValue = TypeVar("StateValue", bound=Mapping[str, Any])
+SessionValue = TypeVar("SessionValue", bound=Mapping[str, Any])
 
 
 @runtime_checkable
-class StateStore(Protocol[StateValue]):
+class SessionStore(Protocol[SessionValue]):
     """Async key/value store with per-key TTL, generic over its value schema.
 
     Backs short-lived auth state. Each store is parameterized by the typed
@@ -14,11 +14,11 @@ class StateStore(Protocol[StateValue]):
     it out, so several stores can share one Redis instance without colliding.
     """
 
-    async def get(self, key: str) -> Optional[StateValue]: ...
+    async def get(self, key: str) -> Optional[SessionValue]: ...
 
-    async def set(self, key: str, value: StateValue, ttl_seconds: Optional[int] = None) -> None: ...
+    async def set(self, key: str, value: SessionValue, ttl_seconds: Optional[int] = None) -> None: ...
 
-    async def pop(self, key: str) -> Optional[StateValue]: ...
+    async def pop(self, key: str) -> Optional[SessionValue]: ...
 
     async def delete(self, key: str) -> None: ...
 

litellm/proxy/auth_v2/sessions/factory.py

@@ -1,79 +0,0 @@
-from __future__ import annotations
-
-import logging
-import os
-from typing import TYPE_CHECKING, Optional, cast
-
-from litellm._redis import get_redis_async_client
-
-from litellm.proxy.auth_v2.sessions.base import StateStore, StateValue
-from litellm.proxy.auth_v2.sessions.memory import InMemoryStateStore
-from litellm.proxy.auth_v2.sessions.redis import RedisStateStore
-
-if TYPE_CHECKING:
-    from redis.asyncio import Redis
-
-logger = logging.getLogger("litellm.proxy.auth_v2.sessions")
-
-_REDIS_ENV_SIGNALS = (
-    "REDIS_URL",
-    "REDIS_HOST",
-    "REDIS_CLUSTER_NODES",
-    "REDIS_SENTINEL_NODES",
-)
-
-
-async def _reachable(client: "Redis") -> bool:
-    try:
-        return bool(await client.ping())
-    except Exception:
-        return False
-
-
-def _default_redis_client() -> Optional["Redis"]:
-    if not any(os.getenv(signal) for signal in _REDIS_ENV_SIGNALS):
-        return None
-    try:
-        return cast("Redis", get_redis_async_client())
-    except Exception:
-        logger.warning("auth_v2 state layer could not build a Redis client", exc_info=True)
-        return None
-
-
-class StateBackend:
-    """Hands out namespaced state stores backed by Redis when reachable, else memory.
-
-    The Redis-vs-memory choice is made once, at ``connect`` time, and held for the
-    backend's lifetime. We deliberately do not fail over per operation: silently
-    moving a live session from Redis to a local dict would strand it on one worker
-    and lose it the moment another worker serves the next request.
-
-    Inject the client for tests or to share the proxy's existing connection; the
-    default builder only fires when Redis is configured via the environment.
-    """
-
-    def __init__(self, redis_client: Optional["Redis"]) -> None:
-        self._redis = redis_client
-
-    @classmethod
-    async def connect(cls, redis_client: Optional["Redis"] = None) -> "StateBackend":
-        client = redis_client if redis_client is not None else _default_redis_client()
-        if client is not None and await _reachable(client):
-            logger.info("auth_v2 state layer using Redis backend")
-            return cls(client)
-        logger.info("auth_v2 state layer using in-memory backend")
-        return cls(None)
-
-    @property
-    def using_redis(self) -> bool:
-        return self._redis is not None
-
-    def store(self, namespace: str, *, default_ttl: int) -> StateStore[StateValue]:
-        """Return a typed store for ``namespace``.
-
-        The value schema is taken from the call site's annotation, e.g.
-        ``sessions: StateStore[SessionState] = backend.store("sessions", default_ttl=3600)``.
-        """
-        if self._redis is not None:
-            return RedisStateStore(self._redis, namespace, default_ttl)
-        return InMemoryStateStore(namespace, default_ttl)

litellm/proxy/auth_v2/sessions/memory.py

@@ -3,17 +3,17 @@ from __future__ import annotations
 import time
 from typing import Dict, Generic, Optional, Tuple
 
-from litellm.proxy.auth_v2.sessions.base import StateValue
+from litellm.proxy.auth_v2.sessions.base import SessionValue
 
 
-class InMemoryStateStore(Generic[StateValue]):
+class InMemorySessionStore(Generic[SessionValue]):
     """Process-local fallback when Redis is unavailable. Single-process only."""
 
     def __init__(self, namespace: str, default_ttl: int, max_size: int = 10000) -> None:
         self._namespace = namespace
         self._default_ttl = default_ttl
         self._max_size = max_size
-        self._entries: Dict[str, Tuple[float, Optional[StateValue]]] = {}
+        self._entries: Dict[str, Tuple[float, Optional[SessionValue]]] = {}
 
     def _key(self, key: str) -> str:
         return f"{self._namespace}:{key}"
@@ -21,7 +21,7 @@ class InMemoryStateStore(Generic[StateValue]):
     def _expiry(self, ttl_seconds: Optional[int]) -> float:
         return time.time() + (self._default_ttl if ttl_seconds is None else ttl_seconds)
 
-    def _live(self, key: str, now: float) -> Optional[Tuple[float, Optional[StateValue]]]:
+    def _live(self, key: str, now: float) -> Optional[Tuple[float, Optional[SessionValue]]]:
         entry = self._entries.get(key)
         if entry is None:
             return None
@@ -30,15 +30,15 @@ class InMemoryStateStore(Generic[StateValue]):
             return None
         return entry
 
-    async def get(self, key: str) -> Optional[StateValue]:
+    async def get(self, key: str) -> Optional[SessionValue]:
         entry = self._live(self._key(key), time.time())
         return entry[1] if entry is not None else None
 
-    async def set(self, key: str, value: StateValue, ttl_seconds: Optional[int] = None) -> None:
+    async def set(self, key: str, value: SessionValue, ttl_seconds: Optional[int] = None) -> None:
         self._evict(time.time())
         self._entries[self._key(key)] = (self._expiry(ttl_seconds), value)
 
-    async def pop(self, key: str) -> Optional[StateValue]:
+    async def pop(self, key: str) -> Optional[SessionValue]:
         entry = self._entries.pop(self._key(key), None)
         if entry is None or entry[0] < time.time():
             return None

litellm/proxy/auth_v2/sessions/redis.py

@@ -3,13 +3,13 @@ from __future__ import annotations
 import json
 from typing import TYPE_CHECKING, Generic, Optional, cast
 
-from litellm.proxy.auth_v2.sessions.base import StateValue
+from litellm.proxy.auth_v2.sessions.base import SessionValue
 
 if TYPE_CHECKING:
     from redis.asyncio import Redis
 
 
-class RedisStateStore(Generic[StateValue]):
+class RedisSessionStore(Generic[SessionValue]):
     """Redis-backed store. Shared across workers; Redis enforces the TTL."""
 
     def __init__(self, client: "Redis", namespace: str, default_ttl: int) -> None:
@@ -23,16 +23,16 @@ class RedisStateStore(Generic[StateValue]):
     def _ttl(self, ttl_seconds: Optional[int]) -> int:
         return self._default_ttl if ttl_seconds is None else ttl_seconds
 
-    async def get(self, key: str) -> Optional[StateValue]:
+    async def get(self, key: str) -> Optional[SessionValue]:
         raw = await self._client.get(self._key(key))
-        return cast(StateValue, json.loads(raw)) if raw is not None else None
+        return cast(SessionValue, json.loads(raw)) if raw is not None else None
 
-    async def set(self, key: str, value: StateValue, ttl_seconds: Optional[int] = None) -> None:
+    async def set(self, key: str, value: SessionValue, ttl_seconds: Optional[int] = None) -> None:
         await self._client.set(self._key(key), json.dumps(value), ex=self._ttl(ttl_seconds))
 
-    async def pop(self, key: str) -> Optional[StateValue]:
+    async def pop(self, key: str) -> Optional[SessionValue]:
         raw = await self._client.getdel(self._key(key))
-        return cast(StateValue, json.loads(raw)) if raw is not None else None
+        return cast(SessionValue, json.loads(raw)) if raw is not None else None
 
     async def delete(self, key: str) -> None:
         await self._client.delete(self._key(key))