mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-17 04:34:25 +00:00
fix(proxy): share CLI SSO login sessions across workers without enable_redis_auth_cache
This commit is contained in:
Devin AI
ryan-crabbe-berri
parent
06e8013e6c
commit
37a4db0743
4 changed files with 104 additions and 31 deletions
|
|
@ -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 = {
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
@ -1955,6 +1956,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
|
||||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -2231,7 +2231,10 @@ class TestCLIKeyRegenerationFlow:
|
|||
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 @@ class TestCLIKeyRegenerationFlow:
|
|||
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 @@ 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,
|
||||
|
|
@ -2544,6 +2551,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,
|
||||
|
|
@ -2582,6 +2590,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>",
|
||||
|
|
@ -2618,7 +2627,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"
|
||||
|
|
@ -2651,7 +2663,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"
|
||||
|
|
@ -2709,6 +2724,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>",
|
||||
|
|
@ -2777,7 +2793,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,
|
||||
|
|
@ -2816,7 +2835,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)
|
||||
|
||||
|
|
@ -2843,7 +2865,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,
|
||||
|
|
@ -3021,6 +3046,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",
|
||||
|
|
@ -3095,6 +3121,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",
|
||||
|
|
@ -3157,6 +3184,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",
|
||||
|
|
@ -6281,6 +6309,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(
|
||||
|
|
@ -6328,6 +6357,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",
|
||||
|
|
@ -6402,6 +6432,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",
|
||||
|
|
@ -6451,7 +6482,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,
|
||||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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})"
|
||||
)
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue