diff --git a/backend/auth/routers/oidc.py b/backend/auth/routers/oidc.py index a489012ce4c..f7a51fee205 100644 --- a/backend/auth/routers/oidc.py +++ b/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 diff --git a/backend/auth/routers/saml.py b/backend/auth/routers/saml.py index a4b1911dfb1..f04b6dcb7bb 100644 --- a/backend/auth/routers/saml.py +++ b/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, diff --git a/litellm/proxy/auth_v2/authenticators/session.py b/litellm/proxy/auth_v2/authenticators/session.py index 1d9aa7e3c14..907e826409b 100644 --- a/litellm/proxy/auth_v2/authenticators/session.py +++ b/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 diff --git a/litellm/proxy/auth_v2/network.py b/litellm/proxy/auth_v2/network.py index e92ce794174..2ae43dd2292 100644 --- a/litellm/proxy/auth_v2/network.py +++ b/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, diff --git a/litellm/proxy/auth_v2/security.py b/litellm/proxy/auth_v2/security.py index be667a9f398..87f15da4251 100644 --- a/litellm/proxy/auth_v2/security.py +++ b/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 = ( diff --git a/litellm/proxy/auth_v2/sessions/__init__.py b/litellm/proxy/auth_v2/sessions/__init__.py index 7228c9ee3dd..29613496416 100644 --- a/litellm/proxy/auth_v2/sessions/__init__.py +++ b/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", ] diff --git a/litellm/proxy/auth_v2/sessions/base.py b/litellm/proxy/auth_v2/sessions/base.py index a7891997076..e82353dc5a7 100644 --- a/litellm/proxy/auth_v2/sessions/base.py +++ b/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: ... diff --git a/litellm/proxy/auth_v2/sessions/factory.py b/litellm/proxy/auth_v2/sessions/factory.py deleted file mode 100644 index 6b8450f7989..00000000000 --- a/litellm/proxy/auth_v2/sessions/factory.py +++ /dev/null @@ -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) diff --git a/litellm/proxy/auth_v2/sessions/memory.py b/litellm/proxy/auth_v2/sessions/memory.py index ccccfefab5e..ce19b07c22e 100644 --- a/litellm/proxy/auth_v2/sessions/memory.py +++ b/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 diff --git a/litellm/proxy/auth_v2/sessions/redis.py b/litellm/proxy/auth_v2/sessions/redis.py index 489060904b9..7242c89ebda 100644 --- a/litellm/proxy/auth_v2/sessions/redis.py +++ b/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)) diff --git a/litellm/proxy/auth_v2/sessions/schemas.py b/litellm/proxy/auth_v2/sessions/types.py similarity index 100% rename from litellm/proxy/auth_v2/sessions/schemas.py rename to litellm/proxy/auth_v2/sessions/types.py