From 37a4db0743ecb90d697c4eaae6e777c2b1e59bf5 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Tue, 14 Jul 2026 19:19:12 +0000 Subject: [PATCH 1/4] fix(proxy): share CLI SSO login sessions across workers without enable_redis_auth_cache --- litellm/proxy/management_endpoints/ui_sso.py | 29 ++++++----- litellm/proxy/proxy_server.py | 15 +++++- .../proxy/management_endpoints/test_ui_sso.py | 51 ++++++++++++++++--- .../proxy/test_redis_auth_cache_flag.py | 40 ++++++++++++--- 4 files changed, 104 insertions(+), 31 deletions(-) diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py index 0475566192e2..2794d14436b9 100644 --- a/litellm/proxy/management_endpoints/ui_sso.py +++ b/litellm/proxy/management_endpoints/ui_sso.py @@ -588,11 +588,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)), ) @@ -607,7 +607,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 = ( ( @@ -639,9 +639,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") @@ -665,7 +665,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) @@ -856,10 +856,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, ) @@ -907,7 +907,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( @@ -1941,6 +1941,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, ): @@ -1987,7 +1988,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)}" @@ -2017,13 +2018,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) @@ -2063,6 +2065,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, ) @@ -2097,10 +2100,10 @@ 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") @@ -2207,7 +2210,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 6661474d215b..515d179fe265 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -224,6 +224,7 @@ def generate_feedback_box(): 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, @@ -1955,6 +1956,7 @@ async def root_redirect(): 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 @@ -3679,13 +3681,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 92d1b870d757..14f489a71551 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py +++ b/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py @@ -2231,7 +2231,10 @@ async def test_cli_sso_start_creates_bound_flow(self): mock_cache = MagicMock() 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-") @@ -2262,7 +2265,10 @@ async def test_cli_sso_start_rate_limits_by_client_ip(self): mock_cache = MagicMock() 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) @@ -2358,6 +2364,7 @@ async def drive(enabled: bool): 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, @@ -2544,6 +2551,7 @@ async def test_cli_sso_callback_prefills_user_code_on_verify_page(self): ), 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, @@ -2582,6 +2590,7 @@ async def test_cli_sso_complete_verifies_user_code(self): 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", @@ -2618,7 +2627,10 @@ async def test_cli_sso_complete_requires_callback_token(self): "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" @@ -2651,7 +2663,10 @@ async def test_cli_sso_complete_waits_for_callback_before_token_checks(self): "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" @@ -2709,6 +2724,7 @@ async def test_cli_sso_callback_stores_session(self): ), 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", @@ -2777,7 +2793,10 @@ async def test_cli_poll_key_returns_teams_for_selection(self): "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, @@ -2816,7 +2835,10 @@ async def test_cli_poll_key_requires_poll_secret(self): }, } - 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) @@ -2843,7 +2865,10 @@ async def test_cli_poll_key_waits_for_user_code_verification(self): }, } - 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, @@ -3021,6 +3046,7 @@ async def test_cli_poll_key_generates_jwt_with_team(self): 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", @@ -3095,6 +3121,7 @@ async def test_cli_poll_key_does_not_cap_session_when_user_has_budget(self): 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", @@ -3157,6 +3184,7 @@ async def test_cli_poll_key_caps_session_when_user_and_team_have_no_budget(self) 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", @@ -6281,6 +6309,7 @@ async def test_cli_sso_callback_passes_user_defined_values_for_new_users(self): ), 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( @@ -6328,6 +6357,7 @@ async def test_cli_sso_callback_rejects_restricted_sso_group(self): ) 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", @@ -6402,6 +6432,7 @@ async def test_cli_sso_callback_persists_attribution_metadata(self, monkeypatch) ), 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", @@ -6451,7 +6482,10 @@ async def test_cli_poll_key_returns_attribution_metadata(self, monkeypatch): "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, @@ -7314,6 +7348,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 d0cb5ec54657..849d5494c6e7 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 @@ def test_flag_true_attaches_redis_to_user_api_key_cache(self): 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 @@ def test_flag_false_leaves_user_api_key_cache_in_memory_only(self): 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 @@ def test_flag_absent_leaves_user_api_key_cache_in_memory_only(self): 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 @@ def test_spend_counter_cache_always_gets_redis_regardless_of_flag(self): 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 @@ def test_flag_false_spend_gets_redis_but_user_cache_does_not(self): 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})" + ) From d0a9ccac57dd15bc0dbb8ad0ad9883369be75046 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Tue, 14 Jul 2026 19:53:03 +0000 Subject: [PATCH 2/4] 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. --- litellm/proxy/management_endpoints/ui_sso.py | 17 ++-- .../proxy/management_endpoints/test_ui_sso.py | 95 +++++++++++++------ 2 files changed, 78 insertions(+), 34 deletions(-) diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py index 2794d14436b9..b920e8cea41a 100644 --- a/litellm/proxy/management_endpoints/ui_sso.py +++ b/litellm/proxy/management_endpoints/ui_sso.py @@ -253,7 +253,11 @@ 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 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, " @@ -273,11 +277,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=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: 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 14f489a71551..a088b1d31e9d 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py +++ b/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py @@ -2216,6 +2216,45 @@ def test_cli_sso_flow_lookup_tells_legacy_clients_to_upgrade(self): assert "session not found or expired" in expired_exc.value.detail assert "enable_redis_auth_cache" 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=fresh_flow, ttl=CLI_SSO_SESSION_TTL_SECONDS + ) + cache.set_cache.assert_not_called() + @pytest.mark.asyncio async def test_cli_sso_start_creates_bound_flow(self): """Test CLI SSO start creates a polling secret bound flow""" @@ -2228,7 +2267,7 @@ async def test_cli_sso_start_creates_bound_flow(self): 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 ( @@ -2262,7 +2301,7 @@ async def test_cli_sso_start_rate_limits_by_client_ip(self): 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 ( @@ -2287,7 +2326,7 @@ async def test_cli_sso_start_returns_verification_uri_complete_when_enabled(self 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 ( @@ -2321,7 +2360,7 @@ async def test_cli_sso_start_omits_verification_uri_complete_by_default(self): 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 ( @@ -2355,7 +2394,7 @@ async def test_google_login_only_threads_user_code_when_enabled(self): 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): @@ -2532,7 +2571,7 @@ async def test_cli_sso_callback_prefills_user_code_on_verify_page(self): ) 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", @@ -2576,7 +2615,7 @@ async def test_cli_sso_complete_verifies_user_code(self): 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( @@ -2615,7 +2654,7 @@ async def test_cli_sso_complete_requires_callback_token(self): 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( @@ -2652,7 +2691,7 @@ async def test_cli_sso_complete_waits_for_callback_before_token_checks(self): 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( @@ -2702,7 +2741,7 @@ async def test_cli_sso_callback_stores_session(self): 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", @@ -2785,7 +2824,7 @@ async def test_cli_poll_key_returns_teams_for_selection(self): } # 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, @@ -2822,7 +2861,7 @@ async def test_cli_poll_key_requires_poll_secret(self): 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, @@ -2852,7 +2891,7 @@ async def test_cli_poll_key_waits_for_user_code_verification(self): 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, @@ -3034,7 +3073,7 @@ async def test_cli_poll_key_generates_jwt_with_team(self): ) # 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, @@ -3110,7 +3149,7 @@ async def test_cli_poll_key_does_not_cap_session_when_user_has_budget(self): 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, @@ -3173,7 +3212,7 @@ async def test_cli_poll_key_caps_session_when_user_and_team_have_no_budget(self) 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, @@ -4125,7 +4164,7 @@ async def test_prepare_token_exchange_parameters_with_pkce(self): 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} @@ -4176,7 +4215,7 @@ async def test_get_generic_sso_redirect_response_with_pkce(self): 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() @@ -4700,7 +4739,7 @@ async def test_pkce_cache_miss_raises_proxy_exception(self): 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) @@ -4826,7 +4865,7 @@ async def test_pkce_cache_miss_unexpected_format_raises_proxy_exception(self): 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() @@ -4868,7 +4907,7 @@ async def test_pkce_cache_miss_non_strict_logs_warning_and_continues(self, caplo 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) @@ -4956,7 +4995,7 @@ async def test_pkce_cache_miss_unexpected_format_non_strict_logs_warning( 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() @@ -5008,7 +5047,7 @@ async def test_pkce_legacy_string_cache_format_backward_compat(self): 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) @@ -6292,7 +6331,7 @@ async def test_cli_sso_callback_passes_user_defined_values_for_new_users(self): 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", @@ -6334,7 +6373,7 @@ async def test_cli_sso_callback_rejects_restricted_sso_group(self): 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", @@ -6404,7 +6443,7 @@ async def test_cli_sso_callback_persists_attribution_metadata(self, monkeypatch) "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", @@ -6474,7 +6513,7 @@ async def test_cli_poll_key_returns_attribution_metadata(self, monkeypatch): "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, @@ -7336,7 +7375,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, From 894f1a6661b0c09c33e1a8af59c78f1bbab65508 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Wed, 15 Jul 2026 16:19:11 -0700 Subject: [PATCH 3/4] 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. --- litellm/proxy/management_endpoints/ui_sso.py | 8 ++- .../proxy/management_endpoints/test_ui_sso.py | 50 ++++++++++++++++++- 2 files changed, 56 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py index b920e8cea41a..e426e20d1e2d 100644 --- a/litellm/proxy/management_endpoints/ui_sso.py +++ b/litellm/proxy/management_endpoints/ui_sso.py @@ -12,6 +12,7 @@ import base64 import hashlib import inspect +import json import os import re import secrets @@ -258,6 +259,11 @@ def _get_cli_sso_flow_or_raise(login_id: Optional[str], cache: DualCache) -> dic 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, " @@ -280,7 +286,7 @@ def _set_cli_sso_flow(login_id: str, cache: DualCache, flow: dict) -> None: 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=flow, ttl=CLI_SSO_SESSION_TTL_SECONDS) + 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) 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 a088b1d31e9d..5aa663b8e013 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py +++ b/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py @@ -2251,10 +2251,58 @@ def test_cli_sso_flow_is_redis_authoritative_when_redis_attached(self): _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=fresh_flow, ttl=CLI_SSO_SESSION_TTL_SECONDS + 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""" From 25c27b6c0e033bf67a0f2c619fbbb27668d24fb0 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Mon, 20 Jul 2026 14:12:30 -0700 Subject: [PATCH 4/4] 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 --- litellm/proxy/management_endpoints/ui_sso.py | 4 ++-- tests/test_litellm/proxy/management_endpoints/test_ui_sso.py | 3 ++- 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py index 06399ede886d..d51ff06b7a58 100644 --- a/litellm/proxy/management_endpoints/ui_sso.py +++ b/litellm/proxy/management_endpoints/ui_sso.py @@ -267,7 +267,7 @@ def _get_cli_sso_flow_or_raise(login_id: Optional[str], cache: DualCache) -> dic 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( @@ -275,7 +275,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." ), ) 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 6e76a82abc0d..fdf593ceb23c 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,8 @@ def test_cli_sso_flow_lookup_tells_legacy_clients_to_upgrade(self): _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): """