scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-08-28 05:25:59 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 48

fix(proxy): share CLI SSO login sessions across workers without enable_redis_auth_cache (#33261)

Browse source 
* fix(proxy): share CLI SSO login sessions across workers without enable_redis_auth_cache

* fix(proxy): make CLI SSO flow state redis-authoritative across workers

The CLI SSO flow is stored in a DualCache whose get_cache is memory-first, so
the worker that served /sso/cli/start keeps serving its stale in-memory flow and
never observes the sso_complete/session_data update another worker writes during
the OAuth callback. Attaching Redis alone is not enough; poll on the original
worker returns pending forever.

Read and write the flow directly through the attached Redis backend when present
so every worker sees the same authoritative state, falling back to the in-memory
DualCache only when no Redis is configured.

* fix(proxy): serialize CLI SSO flow as JSON for the redis round trip

RedisCache stores values via str(value) and parses reads with
json.loads then ast.literal_eval. The completed flow contains a
LitellmUserRoles enum in session_data.user_role, whose repr is not a
parseable literal, so any worker reading the completed flow from redis
raised SyntaxError and returned 400 "CLI login session not found".
Writing the flow as json.dumps makes the round trip lossless (the enum
is a str subclass) and fails loudly at write time if a non-serializable
value is ever added to the flow.

* fix(proxy): point CLI SSO session-not-found hint at configuring Redis

The error message and warning still told users to set enable_redis_auth_cache,
but the CLI SSO session cache now gets Redis unconditionally whenever one is
configured, so that flag no longer affects CLI login

---------

Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
Co-authored-by: ryan-crabbe-berri <ryan@berri.ai>
This commit is contained in:
devin-ai-integration[bot] 2026-07-22 10:29:34 -07:00 • committed by GitHub
parent fa6b209165
commit 17a83aa896
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 240 additions and 68 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

56
litellm/proxy/management_endpoints/ui_sso.py
View file

@ -12,6 +12,7 @@ import asyncio
import base64
import hashlib
import inspect
import json
import os
import re
import secrets
@ -258,11 +259,20 @@ def _get_cli_sso_flow_or_raise(login_id: Optional[str], cache: DualCache) -> dic
raise HTTPException(status_code=400, detail="Invalid CLI login session id")
cache_key = _get_cli_sso_flow_cache_key(cast(str, login_id))
flow = cache.get_cache(key=cache_key)
redis_cache = cache.redis_cache
if redis_cache is not None:
flow = redis_cache.get_cache(key=cache_key)
else:
flow = cache.get_cache(key=cache_key)
if isinstance(flow, str):
try:
flow = json.loads(flow)
except ValueError:
flow = None
if not isinstance(flow, dict) or "poll_secret_hash" not in flow:
verbose_proxy_logger.warning(
"CLI SSO login session not found in cache for login_id=%s. If the proxy runs multiple replicas, "
"a shared Redis cache (enable_redis_auth_cache: true) is required for CLI login to work.",
"a shared Redis cache is required for CLI login to work.",
login_id,
)
raise HTTPException(
@ -270,7 +280,7 @@ def _get_cli_sso_flow_or_raise(login_id: Optional[str], cache: DualCache) -> dic
detail=(
"CLI login session not found or expired. Run `litellm-proxy login` again. "
"If this happens immediately after starting a login, the proxy is likely running multiple "
"replicas without a shared cache; configure Redis with `enable_redis_auth_cache: true` "
"replicas without a shared cache; configure a Redis cache "
"so every replica can see the login session."
),
)
@ -278,11 +288,12 @@ def _get_cli_sso_flow_or_raise(login_id: Optional[str], cache: DualCache) -> dic
def _set_cli_sso_flow(login_id: str, cache: DualCache, flow: dict) -> None:
cache.set_cache(
key=_get_cli_sso_flow_cache_key(login_id),
value=flow,
ttl=CLI_SSO_SESSION_TTL_SECONDS,
)
cache_key = _get_cli_sso_flow_cache_key(login_id)
redis_cache = cache.redis_cache
if redis_cache is not None:
redis_cache.set_cache(key=cache_key, value=json.dumps(flow), ttl=CLI_SSO_SESSION_TTL_SECONDS)
else:
cache.set_cache(key=cache_key, value=flow, ttl=CLI_SSO_SESSION_TTL_SECONDS)
def _verify_cli_sso_poll_secret(flow: dict, poll_secret: Optional[str]) -> bool:
@ -593,11 +604,11 @@ def _render_cli_sso_verification_page(
@router.post("/sso/cli/start", tags=["experimental"], include_in_schema=False)
async def cli_sso_start(request: Request):
from litellm.proxy.proxy_server import general_settings, user_api_key_cache
from litellm.proxy.proxy_server import cli_sso_session_cache, general_settings
_check_cli_sso_start_rate_limit(
request=request,
cache=user_api_key_cache,
cache=cli_sso_session_cache,
use_x_forwarded_for=bool((general_settings or {}).get("use_x_forwarded_for", False)),
)
@ -612,7 +623,7 @@ async def cli_sso_start(request: Request):
"user_code_verified": False,
"session_data": None,
}
_set_cli_sso_flow(login_id=login_id, cache=user_api_key_cache, flow=flow)
_set_cli_sso_flow(login_id=login_id, cache=cli_sso_session_cache, flow=flow)
verification_uri_complete: str | None = (
(
@ -644,9 +655,9 @@ async def cli_sso_complete(request: Request, login_id: str):
from litellm.proxy.common_utils.html_forms.cli_sso_success import (
render_cli_sso_success_page,
)
from litellm.proxy.proxy_server import user_api_key_cache
from litellm.proxy.proxy_server import cli_sso_session_cache
flow = _get_cli_sso_flow_or_raise(login_id=login_id, cache=user_api_key_cache)
flow = _get_cli_sso_flow_or_raise(login_id=login_id, cache=cli_sso_session_cache)
if not flow.get("sso_complete") or not flow.get("session_data"):
raise HTTPException(status_code=400, detail="CLI login is not ready")
@ -670,7 +681,7 @@ async def cli_sso_complete(request: Request, login_id: str):
raise HTTPException(status_code=400, detail="Invalid verification code")
flow["user_code_verified"] = True
_set_cli_sso_flow(login_id=login_id, cache=user_api_key_cache, flow=flow)
_set_cli_sso_flow(login_id=login_id, cache=cli_sso_session_cache, flow=flow)
html_content = render_cli_sso_success_page()
return HTMLResponse(content=html_content, status_code=200)
@ -861,10 +872,10 @@ async def google_login(
Example:
"""
from litellm.proxy.proxy_server import (
cli_sso_session_cache,
general_settings,
premium_user,
prisma_client,
user_api_key_cache,
user_custom_ui_sso_sign_in_handler,
)
@ -912,7 +923,7 @@ async def google_login(
)
if source == LITELLM_CLI_SOURCE_IDENTIFIER:
_get_cli_sso_flow_or_raise(login_id=key, cache=user_api_key_cache)
_get_cli_sso_flow_or_raise(login_id=key, cache=cli_sso_session_cache)
# Store CLI login handle in state for OAuth flow
cli_state: Optional[str] = SSOAuthenticationHandler._get_cli_state(
@ -1957,6 +1968,7 @@ async def _complete_cli_sso_callback_session(
user_defined_values: Optional[SSOUserDefinedValues],
prisma_client: PrismaClient,
user_api_key_cache: UserApiKeyCache,
cli_sso_session_cache: DualCache,
proxy_logging_obj: ProxyLogging,
prefill_user_code: str | None = None,
sso_assertion: SSOIdentityAssertion | None = None,
@ -2006,7 +2018,7 @@ async def _complete_cli_sso_callback_session(
flow["sso_complete"] = True
browser_complete_token = secrets.token_urlsafe(32)
flow["browser_complete_token_hash"] = _hash_cli_sso_secret(browser_complete_token)
_set_cli_sso_flow(login_id=key, cache=user_api_key_cache, flow=flow)
_set_cli_sso_flow(login_id=key, cache=cli_sso_session_cache, flow=flow)
verbose_proxy_logger.info(
f"Stored CLI SSO session for user: {user_info.user_id}, teams: {teams}, num_teams: {len(teams)}"
@ -2037,13 +2049,14 @@ async def cli_sso_callback(
verbose_proxy_logger.info("CLI SSO callback")
from litellm.proxy.proxy_server import (
cli_sso_session_cache,
general_settings,
prisma_client,
proxy_logging_obj,
user_api_key_cache,
)
flow = _get_cli_sso_flow_or_raise(login_id=key, cache=user_api_key_cache)
flow = _get_cli_sso_flow_or_raise(login_id=key, cache=cli_sso_session_cache)
if prisma_client is None:
raise HTTPException(status_code=500, detail=CommonProxyErrors.db_not_connected_error.value)
@ -2083,6 +2096,7 @@ async def cli_sso_callback(
user_defined_values=user_defined_values,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
cli_sso_session_cache=cli_sso_session_cache,
proxy_logging_obj=proxy_logging_obj,
prefill_user_code=prefill_user_code,
sso_assertion=sso_assertion,
@ -2114,10 +2128,10 @@ async def cli_poll_key(
team_id: Optional team ID to assign to the JWT. If provided, must be one of user's teams.
"""
from litellm.proxy.auth.auth_checks import ExperimentalUIJWTToken
from litellm.proxy.proxy_server import user_api_key_cache
from litellm.proxy.proxy_server import cli_sso_session_cache
try:
flow = _get_cli_sso_flow_or_raise(login_id=key_id, cache=user_api_key_cache)
flow = _get_cli_sso_flow_or_raise(login_id=key_id, cache=cli_sso_session_cache)
if not _verify_cli_sso_poll_secret(flow=flow, poll_secret=x_litellm_cli_poll_secret):
raise HTTPException(status_code=403, detail="Invalid CLI polling secret")
@ -2192,7 +2206,7 @@ async def cli_poll_key(
)
# Delete cache entry (single-use)
user_api_key_cache.delete_cache(key=_get_cli_sso_flow_cache_key(key_id))
cli_sso_session_cache.delete_cache(key=_get_cli_sso_flow_cache_key(key_id))
verbose_proxy_logger.info(f"CLI JWT generated for user: {user_id}, team: {team_id}")
poll_response = {

15
litellm/proxy/proxy_server.py
View file

@ -226,6 +226,7 @@ from litellm.constants import (
APSCHEDULER_MAX_INSTANCES,
APSCHEDULER_MISFIRE_GRACE_TIME,
APSCHEDULER_REPLACE_EXISTING,
CLI_SSO_SESSION_TTL_SECONDS,
DAYS_IN_A_MONTH,
DEFAULT_HEALTH_CHECK_INTERVAL,
DEFAULT_MODEL_CREATED_AT_TIME,
@ -1970,6 +1971,7 @@ user_api_key_cache: UserApiKeyCache = UserApiKeyCache(
default_in_memory_ttl=UserAPIKeyCacheTTLEnum.in_memory_cache_ttl.value
)
spend_counter_cache = DualCache(default_in_memory_ttl=UserAPIKeyCacheTTLEnum.in_memory_cache_ttl.value)
cli_sso_session_cache = DualCache(default_in_memory_ttl=CLI_SSO_SESSION_TTL_SECONDS)
model_max_budget_limiter = _PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=user_api_key_cache)
litellm.logging_callback_manager.add_litellm_callback(model_max_budget_limiter)
redis_usage_cache: Optional[RedisCache] = None # redis cache used for tracking spend, tpm/rpm limits
@ -3696,13 +3698,22 @@ def _build_redis_usage_cache_from_environment() -> RedisCache | None:
def _attach_redis_usage_cache(redis_cache: RedisCache, enable_redis_auth_cache: bool) -> None:
"""
Wires an established coordination Redis into the proxy-level caches that
consume it directly: the spend counter cache, the cluster-wide config
cache, and (only when opted in) the virtual-key auth cache.
consume it directly: the spend counter cache, the CLI SSO login-session
cache, the cluster-wide config cache, and (only when opted in) the
virtual-key auth cache.
The CLI SSO login-session cache is always backed by Redis when available so
that the browser SSO flow behind `lite login` survives landing on different
workers; it must not be gated behind enable_redis_auth_cache.
"""
spend_counter_cache.attach_redis_cache(
redis_cache,
default_redis_ttl=litellm.default_redis_ttl,
)
cli_sso_session_cache.attach_redis_cache(
redis_cache,
default_redis_ttl=CLI_SSO_SESSION_TTL_SECONDS,
)
if enable_redis_auth_cache is True:
user_api_key_cache.attach_redis_cache(
redis_cache,

197
tests/test_litellm/proxy/management_endpoints/test_ui_sso.py
View file

@ -2214,7 +2214,95 @@ class TestCLIKeyRegenerationFlow:
_get_cli_sso_flow_or_raise(login_id="cli-test_1234567890", cache=mock_cache)
assert expired_exc.value.status_code == 400
assert "session not found or expired" in expired_exc.value.detail
assert "enable_redis_auth_cache" in expired_exc.value.detail
assert "configure a Redis cache" in expired_exc.value.detail
assert "enable_redis_auth_cache" not in expired_exc.value.detail
def test_cli_sso_flow_is_redis_authoritative_when_redis_attached(self):
"""
When Redis is attached, the CLI SSO flow must be read from and written to
Redis directly, never the in-memory layer. Otherwise the worker that served
/sso/cli/start keeps serving its stale in-memory flow and never sees the
sso_complete/session_data update another worker wrote, which is exactly the
multi-worker failure this fix targets.
"""
from litellm.proxy.management_endpoints.ui_sso import (
CLI_SSO_SESSION_TTL_SECONDS,
_get_cli_sso_flow_cache_key,
_get_cli_sso_flow_or_raise,
_set_cli_sso_flow,
)
login_id = "cli-redis_authoritative_1234567890"
cache_key = _get_cli_sso_flow_cache_key(login_id)
fresh_flow = {"poll_secret_hash": "fresh", "sso_complete": True}
stale_flow = {"poll_secret_hash": "stale", "sso_complete": False}
redis_cache = MagicMock()
redis_cache.get_cache.return_value = fresh_flow
cache = MagicMock()
cache.redis_cache = redis_cache
cache.get_cache.return_value = stale_flow
result = _get_cli_sso_flow_or_raise(login_id=login_id, cache=cache)
assert result == fresh_flow
redis_cache.get_cache.assert_called_once_with(key=cache_key)
cache.get_cache.assert_not_called()
_set_cli_sso_flow(login_id=login_id, cache=cache, flow=fresh_flow)
redis_cache.set_cache.assert_called_once_with(
key=cache_key, value=json.dumps(fresh_flow), ttl=CLI_SSO_SESSION_TTL_SECONDS
)
cache.set_cache.assert_not_called()
def test_cli_sso_flow_with_enum_survives_redis_round_trip(self):
"""
RedisCache stores values via str(value) and reads them back through
json.loads/ast.literal_eval. A raw flow dict containing a Python enum
(session_data.user_role after the SSO callback) produces an unparseable
repr, so every worker reading the completed flow from Redis got a
SyntaxError and returned 400 "session not found". The flow must survive
a real Redis serialization round trip.
"""
from litellm.caching.redis_cache import RedisCache
from litellm.proxy._types import LitellmUserRoles
from litellm.proxy.management_endpoints.ui_sso import (
_get_cli_sso_flow_or_raise,
_set_cli_sso_flow,
)
login_id = "cli-enum_round_trip_1234567890"
completed_flow = {
"poll_secret_hash": "hash",
"sso_complete": True,
"user_code_verified": False,
"session_data": {
"user_id": "user-1",
"user_role": LitellmUserRoles.INTERNAL_USER_VIEW_ONLY,
"models": [],
"teams": ["team-1"],
"team_details": [{"team_id": "team-1", "team_alias": "alias"}],
},
}
redis_store: dict = {}
redis_cache = MagicMock()
redis_cache.set_cache.side_effect = lambda key, value, ttl: redis_store.__setitem__(
key, str(value).encode("utf-8")
)
redis_cache.get_cache.side_effect = lambda key: RedisCache._get_cache_logic(
MagicMock(), redis_store.get(key)
)
cache = MagicMock()
cache.redis_cache = redis_cache
_set_cli_sso_flow(login_id=login_id, cache=cache, flow=completed_flow)
flow = _get_cli_sso_flow_or_raise(login_id=login_id, cache=cache)
assert flow["sso_complete"] is True
assert flow["session_data"]["user_role"] == LitellmUserRoles.INTERNAL_USER_VIEW_ONLY.value
assert flow["session_data"]["team_details"] == [{"team_id": "team-1", "team_alias": "alias"}]
@pytest.mark.asyncio
async def test_cli_sso_start_creates_bound_flow(self):
@ -2228,10 +2316,13 @@ class TestCLIKeyRegenerationFlow:
mock_request = MagicMock(spec=Request)
mock_request.client = SimpleNamespace(host="127.0.0.1")
mock_request.headers = {}
mock_cache = MagicMock()
mock_cache = MagicMock(redis_cache=None)
mock_cache.increment_cache.return_value = 1
with patch("litellm.proxy.proxy_server.user_api_key_cache", mock_cache):
with (
patch("litellm.proxy.proxy_server.user_api_key_cache", mock_cache),
patch("litellm.proxy.proxy_server.cli_sso_session_cache", mock_cache),
):
result = await cli_sso_start(request=mock_request)
assert result["login_id"].startswith("cli-")
@ -2259,10 +2350,13 @@ class TestCLIKeyRegenerationFlow:
mock_request = MagicMock(spec=Request)
mock_request.client = SimpleNamespace(host="127.0.0.1")
mock_request.headers = {}
mock_cache = MagicMock()
mock_cache = MagicMock(redis_cache=None)
mock_cache.increment_cache.return_value = 31
with patch("litellm.proxy.proxy_server.user_api_key_cache", mock_cache):
with (
patch("litellm.proxy.proxy_server.user_api_key_cache", mock_cache),
patch("litellm.proxy.proxy_server.cli_sso_session_cache", mock_cache),
):
with pytest.raises(HTTPException) as exc_info:
await cli_sso_start(request=mock_request)
@ -2281,7 +2375,7 @@ class TestCLIKeyRegenerationFlow:
mock_request.client = SimpleNamespace(host="127.0.0.1")
mock_request.headers = {}
mock_request.base_url = "https://proxy.example.com/"
mock_cache = MagicMock()
mock_cache = MagicMock(redis_cache=None)
mock_cache.increment_cache.return_value = 1
with (
@ -2315,7 +2409,7 @@ class TestCLIKeyRegenerationFlow:
mock_request.client = SimpleNamespace(host="127.0.0.1")
mock_request.headers = {}
mock_request.base_url = "https://proxy.example.com/"
mock_cache = MagicMock()
mock_cache = MagicMock(redis_cache=None)
mock_cache.increment_cache.return_value = 1
with (
@ -2349,7 +2443,7 @@ class TestCLIKeyRegenerationFlow:
mock_request = MagicMock(spec=Request)
mock_request.base_url = "https://proxy.example.com/"
mock_cache = MagicMock()
mock_cache = MagicMock(redis_cache=None)
mock_cache.get_cache.return_value = {"poll_secret_hash": "h"}
async def drive(enabled: bool):
@ -2358,6 +2452,7 @@ class TestCLIKeyRegenerationFlow:
patch("litellm.proxy.proxy_server.premium_user", True),
patch("litellm.proxy.proxy_server.prisma_client", MagicMock()),
patch("litellm.proxy.proxy_server.user_api_key_cache", mock_cache),
patch("litellm.proxy.proxy_server.cli_sso_session_cache", mock_cache),
patch(
"litellm.proxy.proxy_server.user_custom_ui_sso_sign_in_handler",
None,
@ -2525,7 +2620,7 @@ class TestCLIKeyRegenerationFlow:
)
mock_sso_result = {"user_email": "test@example.com", "user_id": "test-user-123"}
mock_cache = MagicMock()
mock_cache = MagicMock(redis_cache=None)
mock_cache.get_cache.return_value = {
"poll_secret_hash": "poll-secret-hash",
"user_code_hash": "user-code-hash",
@ -2544,6 +2639,7 @@ class TestCLIKeyRegenerationFlow:
),
patch("litellm.proxy.proxy_server.prisma_client", MagicMock()),
patch("litellm.proxy.proxy_server.user_api_key_cache", mock_cache),
patch("litellm.proxy.proxy_server.cli_sso_session_cache", mock_cache),
):
result = await cli_sso_callback(
request=mock_request,
@ -2568,7 +2664,7 @@ class TestCLIKeyRegenerationFlow:
mock_request.body = AsyncMock(
return_value=b"user_code=ABCD-EFGH&browser_complete_token=browser-token"
)
mock_cache = MagicMock()
mock_cache = MagicMock(redis_cache=None)
mock_cache.get_cache.return_value = {
"poll_secret_hash": _hash_cli_sso_secret("poll-secret"),
"user_code_hash": _hash_cli_sso_secret(
@ -2582,6 +2678,7 @@ class TestCLIKeyRegenerationFlow:
with (
patch("litellm.proxy.proxy_server.user_api_key_cache", mock_cache),
patch("litellm.proxy.proxy_server.cli_sso_session_cache", mock_cache),
patch(
"litellm.proxy.common_utils.html_forms.cli_sso_success.render_cli_sso_success_page",
return_value="<html>Success</html>",
@ -2606,7 +2703,7 @@ class TestCLIKeyRegenerationFlow:
mock_request = MagicMock(spec=Request)
mock_request.body = AsyncMock(return_value=b"user_code=ABCD-EFGH")
mock_cache = MagicMock()
mock_cache = MagicMock(redis_cache=None)
mock_cache.get_cache.return_value = {
"poll_secret_hash": _hash_cli_sso_secret("poll-secret"),
"user_code_hash": _hash_cli_sso_secret(
@ -2618,7 +2715,10 @@ class TestCLIKeyRegenerationFlow:
"session_data": {"user_id": "test-user-123"},
}
with patch("litellm.proxy.proxy_server.user_api_key_cache", mock_cache):
with (
patch("litellm.proxy.proxy_server.user_api_key_cache", mock_cache),
patch("litellm.proxy.proxy_server.cli_sso_session_cache", mock_cache),
):
with pytest.raises(HTTPException) as exc_info:
await cli_sso_complete(
request=mock_request, login_id="cli-session-4567890"
@ -2640,7 +2740,7 @@ class TestCLIKeyRegenerationFlow:
mock_request.body = AsyncMock(
return_value=b"user_code=ABCD-EFGH&browser_complete_token=browser-token"
)
mock_cache = MagicMock()
mock_cache = MagicMock(redis_cache=None)
mock_cache.get_cache.return_value = {
"poll_secret_hash": _hash_cli_sso_secret("poll-secret"),
"user_code_hash": _hash_cli_sso_secret(
@ -2651,7 +2751,10 @@ class TestCLIKeyRegenerationFlow:
"session_data": None,
}
with patch("litellm.proxy.proxy_server.user_api_key_cache", mock_cache):
with (
patch("litellm.proxy.proxy_server.user_api_key_cache", mock_cache),
patch("litellm.proxy.proxy_server.cli_sso_session_cache", mock_cache),
):
with pytest.raises(HTTPException) as exc_info:
await cli_sso_complete(
request=mock_request, login_id="cli-session-4567890"
@ -2687,7 +2790,7 @@ class TestCLIKeyRegenerationFlow:
mock_sso_result = {"user_email": "test@example.com", "user_id": "test-user-123"}
# Mock cache
mock_cache = MagicMock()
mock_cache = MagicMock(redis_cache=None)
mock_cache.get_cache.return_value = {
"poll_secret_hash": "poll-secret-hash",
"user_code_hash": "user-code-hash",
@ -2709,6 +2812,7 @@ class TestCLIKeyRegenerationFlow:
),
patch("litellm.proxy.proxy_server.prisma_client", MagicMock()),
patch("litellm.proxy.proxy_server.user_api_key_cache", mock_cache),
patch("litellm.proxy.proxy_server.cli_sso_session_cache", mock_cache),
patch(
"litellm.proxy.common_utils.html_forms.cli_sso_success.render_cli_sso_success_page",
return_value="<html>Success</html>",
@ -2769,7 +2873,7 @@ class TestCLIKeyRegenerationFlow:
}
# Mock cache
mock_cache = MagicMock()
mock_cache = MagicMock(redis_cache=None)
mock_cache.get_cache.return_value = {
"poll_secret_hash": _hash_cli_sso_secret("poll-secret"),
"sso_complete": True,
@ -2777,7 +2881,10 @@ class TestCLIKeyRegenerationFlow:
"session_data": session_data,
}
with patch("litellm.proxy.proxy_server.user_api_key_cache", mock_cache):
with (
patch("litellm.proxy.proxy_server.user_api_key_cache", mock_cache),
patch("litellm.proxy.proxy_server.cli_sso_session_cache", mock_cache),
):
# Act - First poll without team_id
result = await cli_poll_key(
key_id=session_key,
@ -2803,7 +2910,7 @@ class TestCLIKeyRegenerationFlow:
cli_poll_key,
)
mock_cache = MagicMock()
mock_cache = MagicMock(redis_cache=None)
mock_cache.get_cache.return_value = {
"poll_secret_hash": _hash_cli_sso_secret("poll-secret"),
"sso_complete": True,
@ -2816,7 +2923,10 @@ class TestCLIKeyRegenerationFlow:
},
}
with patch("litellm.proxy.proxy_server.user_api_key_cache", mock_cache):
with (
patch("litellm.proxy.proxy_server.user_api_key_cache", mock_cache),
patch("litellm.proxy.proxy_server.cli_sso_session_cache", mock_cache),
):
with pytest.raises(HTTPException) as exc_info:
await cli_poll_key(key_id="cli-session-789123", team_id=None)
@ -2830,7 +2940,7 @@ class TestCLIKeyRegenerationFlow:
cli_poll_key,
)
mock_cache = MagicMock()
mock_cache = MagicMock(redis_cache=None)
mock_cache.get_cache.return_value = {
"poll_secret_hash": _hash_cli_sso_secret("poll-secret"),
"sso_complete": True,
@ -2843,7 +2953,10 @@ class TestCLIKeyRegenerationFlow:
},
}
with patch("litellm.proxy.proxy_server.user_api_key_cache", mock_cache):
with (
patch("litellm.proxy.proxy_server.user_api_key_cache", mock_cache),
patch("litellm.proxy.proxy_server.cli_sso_session_cache", mock_cache),
):
result = await cli_poll_key(
key_id="cli-session-789123",
team_id=None,
@ -3011,7 +3124,7 @@ class TestCLIKeyRegenerationFlow:
)
# Mock cache
mock_cache = MagicMock()
mock_cache = MagicMock(redis_cache=None)
mock_cache.get_cache.return_value = {
"poll_secret_hash": _hash_cli_sso_secret("poll-secret"),
"sso_complete": True,
@ -3023,6 +3136,7 @@ class TestCLIKeyRegenerationFlow:
with (
patch("litellm.proxy.proxy_server.user_api_key_cache", mock_cache),
patch("litellm.proxy.proxy_server.cli_sso_session_cache", mock_cache),
patch("litellm.proxy.proxy_server.prisma_client"),
patch(
"litellm.proxy.auth.auth_checks.ExperimentalUIJWTToken.get_cli_jwt_auth_token",
@ -3086,7 +3200,7 @@ class TestCLIKeyRegenerationFlow:
models=["gpt-4"],
max_budget=100.0,
)
mock_cache = MagicMock()
mock_cache = MagicMock(redis_cache=None)
mock_cache.get_cache.return_value = {
"poll_secret_hash": _hash_cli_sso_secret("poll-secret"),
"sso_complete": True,
@ -3097,6 +3211,7 @@ class TestCLIKeyRegenerationFlow:
with (
patch("litellm.proxy.proxy_server.user_api_key_cache", mock_cache),
patch("litellm.proxy.proxy_server.cli_sso_session_cache", mock_cache),
patch("litellm.proxy.proxy_server.prisma_client"),
patch(
"litellm.proxy.auth.auth_checks.ExperimentalUIJWTToken.get_cli_jwt_auth_token",
@ -3142,7 +3257,7 @@ class TestCLIKeyRegenerationFlow:
"models": ["gpt-4"],
"user_email": "unbudgeted@example.com",
}
mock_cache = MagicMock()
mock_cache = MagicMock(redis_cache=None)
mock_cache.get_cache.return_value = {
"poll_secret_hash": _hash_cli_sso_secret("poll-secret"),
"sso_complete": True,
@ -3153,6 +3268,7 @@ class TestCLIKeyRegenerationFlow:
with (
patch("litellm.proxy.proxy_server.user_api_key_cache", mock_cache),
patch("litellm.proxy.proxy_server.cli_sso_session_cache", mock_cache),
patch(
"litellm.proxy.auth.auth_checks.ExperimentalUIJWTToken.get_cli_jwt_auth_token",
return_value=mock_jwt_token,
@ -4082,7 +4198,7 @@ class TestPKCEFunctionality:
mock_request.query_params = {"state": test_state}
# Mock cache with async methods — use dict format (primary path)
mock_cache = MagicMock()
mock_cache = MagicMock(redis_cache=None)
test_code_verifier = "test_code_verifier_abc123xyz"
mock_cache.async_get_cache = AsyncMock(
return_value={"code_verifier": test_code_verifier}
@ -4133,7 +4249,7 @@ class TestPKCEFunctionality:
mock_sso.__exit__ = MagicMock(return_value=False)
test_state = "test456"
mock_cache = MagicMock()
mock_cache = MagicMock(redis_cache=None)
mock_cache.async_set_cache = AsyncMock()
@ -4657,7 +4773,7 @@ class TestPKCEFunctionality:
from litellm.proxy._types import ProxyException
from litellm.proxy.management_endpoints.ui_sso import SSOAuthenticationHandler
mock_cache = MagicMock()
mock_cache = MagicMock(redis_cache=None)
mock_cache.async_get_cache = AsyncMock(return_value=None) # verifier not found
mock_request = MagicMock(spec=Request)
@ -4783,7 +4899,7 @@ class TestPKCEFunctionality:
from litellm.proxy.management_endpoints.ui_sso import SSOAuthenticationHandler
# Cache returns an integer — unexpected format
mock_cache = MagicMock()
mock_cache = MagicMock(redis_cache=None)
mock_cache.async_get_cache = AsyncMock(return_value=12345)
mock_cache.async_delete_cache = AsyncMock()
@ -4825,7 +4941,7 @@ class TestPKCEFunctionality:
from litellm.proxy.management_endpoints.ui_sso import SSOAuthenticationHandler
mock_cache = MagicMock()
mock_cache = MagicMock(redis_cache=None)
mock_cache.async_get_cache = AsyncMock(return_value=None) # verifier not found
mock_request = MagicMock(spec=Request)
@ -4913,7 +5029,7 @@ class TestPKCEFunctionality:
from litellm.proxy.management_endpoints.ui_sso import SSOAuthenticationHandler
# Cache returns an integer — unexpected format
mock_cache = MagicMock()
mock_cache = MagicMock(redis_cache=None)
mock_cache.async_get_cache = AsyncMock(return_value=12345)
mock_cache.async_delete_cache = AsyncMock()
@ -4965,7 +5081,7 @@ class TestPKCEFunctionality:
from litellm.proxy.management_endpoints.ui_sso import SSOAuthenticationHandler
legacy_verifier = "legacy_plain_string_verifier_abc123"
mock_cache = MagicMock()
mock_cache = MagicMock(redis_cache=None)
mock_cache.async_get_cache = AsyncMock(return_value=legacy_verifier)
mock_request = MagicMock(spec=Request)
@ -6249,7 +6365,7 @@ class TestCliSsoAttributionMetadata:
provider="generic",
team_ids=[],
)
mock_cache = MagicMock()
mock_cache = MagicMock(redis_cache=None)
mock_cache.get_cache.return_value = {
"poll_secret_hash": "poll-secret-hash",
"user_code_hash": "user-code-hash",
@ -6266,6 +6382,7 @@ class TestCliSsoAttributionMetadata:
),
patch("litellm.proxy.proxy_server.prisma_client", MagicMock()),
patch("litellm.proxy.proxy_server.user_api_key_cache", mock_cache),
patch("litellm.proxy.proxy_server.cli_sso_session_cache", mock_cache),
patch("litellm.proxy.proxy_server.user_custom_sso", None),
):
await ui_sso.cli_sso_callback(
@ -6290,7 +6407,7 @@ class TestCliSsoAttributionMetadata:
mock_request = MagicMock(spec=Request)
mock_request.base_url = "http://internal-proxy.local/"
mock_cache = MagicMock()
mock_cache = MagicMock(redis_cache=None)
mock_cache.get_cache.return_value = {
"poll_secret_hash": "poll-secret-hash",
"user_code_hash": "user-code-hash",
@ -6313,6 +6430,7 @@ class TestCliSsoAttributionMetadata:
) as get_user_info_mock,
patch("litellm.proxy.proxy_server.prisma_client", MagicMock()),
patch("litellm.proxy.proxy_server.user_api_key_cache", mock_cache),
patch("litellm.proxy.proxy_server.cli_sso_session_cache", mock_cache),
patch("litellm.proxy.proxy_server.user_custom_sso", None),
patch(
"litellm.proxy.proxy_server.general_settings",
@ -6359,7 +6477,7 @@ class TestCliSsoAttributionMetadata:
"user_id": "test-user-123",
"employment_type": "contractor",
}
mock_cache = MagicMock()
mock_cache = MagicMock(redis_cache=None)
mock_cache.get_cache.return_value = {
"poll_secret_hash": "poll-secret-hash",
"user_code_hash": "user-code-hash",
@ -6387,6 +6505,7 @@ class TestCliSsoAttributionMetadata:
),
patch("litellm.proxy.proxy_server.prisma_client", mock_prisma),
patch("litellm.proxy.proxy_server.user_api_key_cache", mock_cache),
patch("litellm.proxy.proxy_server.cli_sso_session_cache", mock_cache),
patch("litellm.proxy.proxy_server.user_custom_sso", None),
patch(
"litellm.proxy.common_utils.html_forms.cli_sso_success.render_cli_sso_success_page",
@ -6428,7 +6547,7 @@ class TestCliSsoAttributionMetadata:
"org": {"cost_center": "CC-42"},
},
}
mock_cache = MagicMock()
mock_cache = MagicMock(redis_cache=None)
mock_cache.get_cache.return_value = {
"poll_secret_hash": _hash_cli_sso_secret("poll-secret"),
"sso_complete": True,
@ -6436,7 +6555,10 @@ class TestCliSsoAttributionMetadata:
"session_data": session_data,
}
with patch("litellm.proxy.proxy_server.user_api_key_cache", mock_cache):
with (
patch("litellm.proxy.proxy_server.user_api_key_cache", mock_cache),
patch("litellm.proxy.proxy_server.cli_sso_session_cache", mock_cache),
):
result = await cli_poll_key(
key_id=session_key,
team_id=None,
@ -7287,7 +7409,7 @@ async def test_cli_poll_key_tolerates_missing_user_row():
"models": ["gpt-4"],
}
mock_cache = MagicMock()
mock_cache = MagicMock(redis_cache=None)
mock_cache.get_cache.return_value = {
"poll_secret_hash": _hash_cli_sso_secret("poll-secret"),
"sso_complete": True,
@ -7299,6 +7421,7 @@ async def test_cli_poll_key_tolerates_missing_user_row():
with (
patch("litellm.proxy.proxy_server.user_api_key_cache", mock_cache),
patch("litellm.proxy.proxy_server.cli_sso_session_cache", mock_cache),
patch("litellm.proxy.proxy_server.prisma_client"),
patch(
"litellm.proxy.auth.auth_checks.ExperimentalUIJWTToken.get_cli_jwt_auth_token",

40
tests/test_litellm/proxy/test_redis_auth_cache_flag.py
View file

@ -54,8 +54,8 @@ def _patched_init_cache(litellm_settings: dict, cache_params: dict):
_FakeRedisCache (passes the isinstance guard in _init_cache).
3. Extracts enable_redis_auth_cache from litellm_settings and passes it
as the second argument to _init_cache (matching production behaviour).
4. Yields (user_api_key_cache, spend_counter_cache) after calling
_init_cache, then restores everything.
4. Yields (user_api_key_cache, spend_counter_cache, cli_sso_session_cache)
after calling _init_cache, then restores everything.
"""
fake_redis = _FakeRedisCache()
@ -64,19 +64,21 @@ def _patched_init_cache(litellm_settings: dict, cache_params: dict):
fresh_user_cache = DualCache()
fresh_spend_cache = DualCache()
fresh_cli_sso_cache = DualCache()
enable_redis_auth_cache = litellm_settings.get("enable_redis_auth_cache", False)
with (
patch.object(ps, "user_api_key_cache", fresh_user_cache),
patch.object(ps, "spend_counter_cache", fresh_spend_cache),
patch.object(ps, "cli_sso_session_cache", fresh_cli_sso_cache),
patch.object(ps, "llm_router", None),
# Cache is locally imported inside _init_cache: patch it at source.
patch("litellm.Cache", return_value=mock_litellm_cache),
):
litellm.cache = None
ps.ProxyConfig()._init_cache(cache_params, enable_redis_auth_cache)
yield fresh_user_cache, fresh_spend_cache
yield fresh_user_cache, fresh_spend_cache, fresh_cli_sso_cache
# ---------------------------------------------------------------------------
@ -90,7 +92,7 @@ class TestRedisAuthCacheFlag:
with _patched_init_cache(
litellm_settings={"enable_redis_auth_cache": True},
cache_params={"type": "redis", "host": "localhost", "port": 6379},
) as (user_cache, _):
) as (user_cache, _, _cli_sso_cache):
assert user_cache.redis_cache is not None, (
"Redis should be attached to user_api_key_cache when "
"enable_redis_auth_cache=True"
@ -101,7 +103,7 @@ class TestRedisAuthCacheFlag:
with _patched_init_cache(
litellm_settings={"enable_redis_auth_cache": False},
cache_params={"type": "redis", "host": "localhost", "port": 6379},
) as (user_cache, _):
) as (user_cache, _, _cli_sso_cache):
assert user_cache.redis_cache is None, (
"user_api_key_cache must remain in-memory-only when "
"enable_redis_auth_cache=False"
@ -112,7 +114,7 @@ class TestRedisAuthCacheFlag:
with _patched_init_cache(
litellm_settings={},
cache_params={"type": "redis", "host": "localhost", "port": 6379},
) as (user_cache, _):
) as (user_cache, _, _cli_sso_cache):
assert user_cache.redis_cache is None, (
"user_api_key_cache must remain in-memory-only when "
"enable_redis_auth_cache is absent from litellm_settings"
@ -129,7 +131,7 @@ class TestRedisAuthCacheFlag:
with _patched_init_cache(
litellm_settings=ls,
cache_params={"type": "redis", "host": "localhost", "port": 6379},
) as (_, spend_cache):
) as (_, spend_cache, _cli_sso_cache):
assert spend_cache.redis_cache is not None, (
f"spend_counter_cache must always get Redis "
f"(enable_redis_auth_cache={flag_value!r})"
@ -140,6 +142,28 @@ class TestRedisAuthCacheFlag:
with _patched_init_cache(
litellm_settings={"enable_redis_auth_cache": False},
cache_params={"type": "redis", "host": "localhost", "port": 6379},
) as (user_cache, spend_cache):
) as (user_cache, spend_cache, _cli_sso_cache):
assert spend_cache.redis_cache is not None
assert user_cache.redis_cache is None
def test_cli_sso_session_cache_always_gets_redis_regardless_of_flag(self):
"""
cli_sso_session_cache must receive Redis regardless of the auth-cache
flag so that `lite login` works on multi-worker deployments without
enable_redis_auth_cache (regression for the CLI SSO "Invalid CLI login
session" bug)
"""
for flag_value in (True, False, None):
ls = (
{"enable_redis_auth_cache": flag_value}
if flag_value is not None
else {}
)
with _patched_init_cache(
litellm_settings=ls,
cache_params={"type": "redis", "host": "localhost", "port": 6379},
) as (_, _, cli_sso_cache):
assert cli_sso_cache.redis_cache is not None, (
f"cli_sso_session_cache must always get Redis "
f"(enable_redis_auth_cache={flag_value!r})"
)