scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-13 23:11:40 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 101

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>
(cherry picked from commit 17a83aa896)
This commit is contained in:
devin-ai-integration[bot] 2026-07-22 10:29:34 -07:00 • committed by Yuneng Jiang
parent 1b562f21ba
commit 4e80c98696
No known key found for this signature in database
4 changed files with 256 additions and 66 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
@ -245,18 +246,28 @@ def _get_cli_sso_flow_or_raise(login_id: Optional[str], cache: DualCache) -> dic
raise HTTPException(status_code=400, detail="Invalid CLI login session")
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:
raise HTTPException(status_code=400, detail="Invalid CLI login session")
return flow
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:
@ -567,11 +578,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)),
)
@ -586,7 +597,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 = (
(
@ -618,9 +629,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")
@ -644,7 +655,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)
@ -835,10 +846,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,
)
@ -886,7 +897,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(
@ -1920,6 +1931,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,
):
@ -1966,7 +1978,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)}"
@ -1996,13 +2008,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)
@ -2042,6 +2055,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,
)
@ -2076,10 +2090,14 @@ async def cli_poll_key(
get_team_object,
get_user_object,
)
from litellm.proxy.proxy_server import prisma_client, user_api_key_cache
from litellm.proxy.proxy_server import (
cli_sso_session_cache,
prisma_client,
user_api_key_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")
@ -2186,7 +2204,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

@ -224,6 +224,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,
@ -1909,6 +1910,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
@ -3632,13 +3634,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,

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

@ -2139,6 +2139,93 @@ class TestCLIKeyRegenerationFlow:
assert not _is_valid_cli_sso_login_id("cli-test\x001234567890")
assert not _is_valid_cli_sso_login_id("sk-test1234567890")
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):
"""Test CLI SSO start creates a polling secret bound flow"""
@ -2151,10 +2238,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-")
@ -2182,10 +2272,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)
@ -2204,7 +2297,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 (
@ -2238,7 +2331,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 (
@ -2272,7 +2365,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):
@ -2281,6 +2374,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,
@ -2448,7 +2542,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",
@ -2467,6 +2561,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,
@ -2491,7 +2586,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(
@ -2505,6 +2600,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>",
@ -2529,7 +2625,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(
@ -2541,7 +2637,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"
@ -2563,7 +2662,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(
@ -2574,7 +2673,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"
@ -2610,7 +2712,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",
@ -2632,6 +2734,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>",
@ -2692,7 +2795,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,
@ -2700,7 +2803,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,
@ -2726,7 +2832,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,
@ -2739,7 +2845,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)
@ -2753,7 +2862,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,
@ -2766,7 +2875,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,
@ -2932,7 +3044,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,
@ -2944,6 +3056,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",
@ -3007,7 +3120,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,
@ -3018,6 +3131,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",
@ -3069,7 +3183,7 @@ class TestCLIKeyRegenerationFlow:
max_budget=None,
)
mock_team = LiteLLM_TeamTableCachedObj(team_id="team-x", max_budget=None)
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,
@ -3080,6 +3194,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",
@ -4020,7 +4135,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}
@ -4071,7 +4186,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()
@ -4595,7 +4710,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)
@ -4721,7 +4836,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()
@ -4763,7 +4878,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)
@ -4851,7 +4966,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()
@ -4903,7 +5018,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)
@ -6187,7 +6302,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",
@ -6204,6 +6319,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(
@ -6228,7 +6344,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",
@ -6251,6 +6367,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",
@ -6297,7 +6414,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",
@ -6325,6 +6442,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",
@ -6366,7 +6484,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,
@ -6374,7 +6492,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,
@ -7225,7 +7346,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,
@ -7237,6 +7358,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",

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

@ -54,8 +54,14 @@ 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.
_init_cache also writes three globals this helper does not patch:
``litellm.cache``, ``ps.redis_usage_cache`` and
``litellm_config_cache.redis_cache``. They are saved and restored here so a
_FakeRedisCache never outlives this context and poisons later test files in
the same pytest session.
"""
fake_redis = _FakeRedisCache()
@ -64,19 +70,30 @@ 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)
prev_litellm_cache = litellm.cache
prev_redis_usage_cache = ps.redis_usage_cache
prev_config_cache_redis = ps.litellm_config_cache.redis_cache
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
try:
ps.ProxyConfig()._init_cache(cache_params, enable_redis_auth_cache)
yield fresh_user_cache, fresh_spend_cache, fresh_cli_sso_cache
finally:
litellm.cache = prev_litellm_cache
ps.redis_usage_cache = prev_redis_usage_cache
ps.litellm_config_cache.redis_cache = prev_config_cache_redis
# ---------------------------------------------------------------------------
@ -90,7 +107,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 +118,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 +129,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 +146,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 +157,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})"
)