diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py index 3c8444ecf26..de988a0140f 100644 --- a/litellm/proxy/management_endpoints/ui_sso.py +++ b/litellm/proxy/management_endpoints/ui_sso.py @@ -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 = { diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 6de3e43fc1a..d4bff81ea6a 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -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, diff --git a/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py b/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py index e1856860c8a..47ceb0c05fa 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py +++ b/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py @@ -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="Success", @@ -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="Success", @@ -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", diff --git a/tests/test_litellm/proxy/test_redis_auth_cache_flag.py b/tests/test_litellm/proxy/test_redis_auth_cache_flag.py index d0cb5ec5465..849d5494c6e 100644 --- a/tests/test_litellm/proxy/test_redis_auth_cache_flag.py +++ b/tests/test_litellm/proxy/test_redis_auth_cache_flag.py @@ -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})" + )