diff --git a/litellm/proxy/client/cli/main.py b/litellm/proxy/client/cli/main.py index 43b64aebd3b..7d34d5a883a 100644 --- a/litellm/proxy/client/cli/main.py +++ b/litellm/proxy/client/cli/main.py @@ -22,6 +22,13 @@ from .commands.users import users from .interface import interactive_shell +def _normalize_base_url(base_url: str) -> str: + """Strip trailing slashes so URL joins like f"{base_url}/sso/cli/start" don't + produce a double slash. Any path prefix (e.g. http://host/gateway) is preserved. + """ + return base_url.rstrip("/") if base_url else base_url + + def print_version(base_url: str, api_key: Optional[str]): """Print CLI and server version info.""" click.echo(f"LiteLLM Proxy CLI Version: {litellm_version}") @@ -49,7 +56,7 @@ def print_version(base_url: str, api_key: Optional[str]): callback=lambda ctx, param, value: ( ( print_version( - ctx.params.get("base_url") or "http://localhost:4000", + _normalize_base_url(ctx.params.get("base_url") or "http://localhost:4000"), ctx.params.get("api_key"), ) or ctx.exit() @@ -76,6 +83,8 @@ def cli(ctx: click.Context, base_url: str, api_key: Optional[str]) -> None: """LiteLLM Proxy CLI - Manage your LiteLLM proxy server""" ctx.ensure_object(dict) + base_url = _normalize_base_url(base_url) + # If no API key provided via flag or environment variable, try to load from saved token. # Pass base_url so we only use the stored key when it was issued for this server. if api_key is None: diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py index 065464aa565..cc8e1e550cb 100644 --- a/litellm/proxy/management_endpoints/ui_sso.py +++ b/litellm/proxy/management_endpoints/ui_sso.py @@ -42,6 +42,7 @@ import litellm from litellm._logging import verbose_proxy_logger from litellm._uuid import uuid from litellm.caching.dual_cache import DualCache +from litellm.caching.redis_cache import RedisCache from litellm.constants import ( CLI_SSO_CLAIM_MAP, CLI_SSO_CLAIM_MAX_SCALAR_LENGTH, @@ -240,19 +241,33 @@ def _check_cli_sso_start_rate_limit( ) -def _get_cli_sso_flow_or_raise(login_id: Optional[str], cache: DualCache) -> dict: +def _get_cli_sso_flow_cache(user_api_key_cache: DualCache) -> Union[DualCache, RedisCache]: + """Resolve the cache backing CLI SSO login sessions. + + The multi-request login flow must be visible to every proxy instance, so it + prefers the shared redis_usage_cache (like the PKCE verifiers) and only falls + back to the per-worker in-memory user_api_key_cache when Redis is not configured. + """ + from litellm.proxy.proxy_server import redis_usage_cache + + if redis_usage_cache is not None: + return redis_usage_cache + return user_api_key_cache + + +async def _get_cli_sso_flow_or_raise(login_id: Optional[str], cache: Union[DualCache, RedisCache]) -> dict: if not _is_valid_cli_sso_login_id(login_id): raise HTTPException(status_code=400, detail="Invalid CLI login session") cache_key = _get_cli_sso_flow_cache_key(cast(str, login_id)) - flow = cache.get_cache(key=cache_key) + flow = await cache.async_get_cache(key=cache_key) if not isinstance(flow, dict) or "poll_secret_hash" not in flow: raise HTTPException(status_code=400, detail="Invalid CLI login session") return flow -def _set_cli_sso_flow(login_id: str, cache: DualCache, flow: dict) -> None: - cache.set_cache( +async def _set_cli_sso_flow(login_id: str, cache: Union[DualCache, RedisCache], flow: dict) -> None: + await cache.async_set_cache( key=_get_cli_sso_flow_cache_key(login_id), value=flow, ttl=CLI_SSO_SESSION_TTL_SECONDS, @@ -586,7 +601,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) + await _set_cli_sso_flow(login_id=login_id, cache=_get_cli_sso_flow_cache(user_api_key_cache), flow=flow) verification_uri_complete: str | None = ( ( @@ -620,7 +635,8 @@ async def cli_sso_complete(request: Request, login_id: str): ) from litellm.proxy.proxy_server import user_api_key_cache - flow = _get_cli_sso_flow_or_raise(login_id=login_id, cache=user_api_key_cache) + flow_cache = _get_cli_sso_flow_cache(user_api_key_cache) + flow = await _get_cli_sso_flow_or_raise(login_id=login_id, cache=flow_cache) if not flow.get("sso_complete") or not flow.get("session_data"): raise HTTPException(status_code=400, detail="CLI login is not ready") @@ -644,7 +660,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) + await _set_cli_sso_flow(login_id=login_id, cache=flow_cache, flow=flow) html_content = render_cli_sso_success_page() return HTMLResponse(content=html_content, status_code=200) @@ -886,7 +902,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) + await _get_cli_sso_flow_or_raise(login_id=key, cache=_get_cli_sso_flow_cache(user_api_key_cache)) # Store CLI login handle in state for OAuth flow cli_state: Optional[str] = SSOAuthenticationHandler._get_cli_state( @@ -1966,7 +1982,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) + await _set_cli_sso_flow(login_id=key, cache=_get_cli_sso_flow_cache(user_api_key_cache), flow=flow) verbose_proxy_logger.info( f"Stored CLI SSO session for user: {user_info.user_id}, teams: {teams}, num_teams: {len(teams)}" @@ -2002,7 +2018,7 @@ async def cli_sso_callback( user_api_key_cache, ) - flow = _get_cli_sso_flow_or_raise(login_id=key, cache=user_api_key_cache) + flow = await _get_cli_sso_flow_or_raise(login_id=key, cache=_get_cli_sso_flow_cache(user_api_key_cache)) if prisma_client is None: raise HTTPException(status_code=500, detail=CommonProxyErrors.db_not_connected_error.value) @@ -2078,8 +2094,9 @@ async def cli_poll_key( ) from litellm.proxy.proxy_server import prisma_client, user_api_key_cache + flow_cache = _get_cli_sso_flow_cache(user_api_key_cache) try: - flow = _get_cli_sso_flow_or_raise(login_id=key_id, cache=user_api_key_cache) + flow = await _get_cli_sso_flow_or_raise(login_id=key_id, cache=flow_cache) if not _verify_cli_sso_poll_secret(flow=flow, poll_secret=x_litellm_cli_poll_secret): raise HTTPException(status_code=403, detail="Invalid CLI polling secret") @@ -2186,7 +2203,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)) + await flow_cache.async_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/tests/test_litellm/proxy/client/cli/test_auth_commands.py b/tests/test_litellm/proxy/client/cli/test_auth_commands.py index 4ee8b502aa2..00f0756aec2 100644 --- a/tests/test_litellm/proxy/client/cli/test_auth_commands.py +++ b/tests/test_litellm/proxy/client/cli/test_auth_commands.py @@ -339,6 +339,55 @@ class TestLoginCommand: # Verify commands were shown mock_show_commands.assert_called_once() + def test_login_normalizes_trailing_slash_base_url(self): + """A trailing slash on --base-url must not produce double-slash URLs. + + Regression: `lite --base-url https://host/ login` used to hit + https://host//sso/cli/start and open https://host//sso/key/generate, + which breaks gateways that 404 on the doubled slash. + """ + from litellm.proxy.client.cli.main import cli + + mock_response = Mock() + mock_response.status_code = 200 + mock_response.json.return_value = { + "status": "ready", + "key": "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.test.jwt", + "user_id": "test-user-123", + "team_id": "team-1", + "teams": ["team-1"], + } + + with ( + patch("webbrowser.open") as mock_browser, + patch( + "requests.post", + return_value=_mock_cli_sso_start_response(login_id="cli-test-uuid-123"), + ) as mock_post, + patch("requests.get", return_value=mock_response) as mock_get, + patch("litellm.proxy.client.cli.commands.auth.save_token") as mock_save, + patch("litellm.proxy.client.cli.commands.auth.get_stored_api_key", return_value=None), + patch("litellm.proxy.client.cli.interface.show_commands"), + ): + result = self.runner.invoke( + cli, ["--base-url", "https://test.example.com/", "login"] + ) + + assert result.exit_code == 0, result.output + + start_url = mock_post.call_args[0][0] + assert start_url == "https://test.example.com/sso/cli/start" + + poll_url = mock_get.call_args[0][0] + assert poll_url.startswith("https://test.example.com/sso/cli/poll/") + assert "//sso" not in poll_url.replace("https://", "") + + browser_url = mock_browser.call_args[0][0] + assert browser_url.startswith("https://test.example.com/sso/key/generate?") + + saved_data = mock_save.call_args[0][0] + assert saved_data["base_url"] == "https://test.example.com" + def test_login_timeout(self): """Test login timeout scenario""" mock_context = Mock() diff --git a/tests/test_litellm/proxy/client/cli/test_global_options.py b/tests/test_litellm/proxy/client/cli/test_global_options.py index 3a19f735c1b..86d93740eb6 100644 --- a/tests/test_litellm/proxy/client/cli/test_global_options.py +++ b/tests/test_litellm/proxy/client/cli/test_global_options.py @@ -13,6 +13,7 @@ sys.path.insert( from litellm._version import version as litellm_version from litellm.proxy.client.cli import cli +from litellm.proxy.client.cli.main import _normalize_base_url @pytest.fixture @@ -20,6 +21,45 @@ def cli_runner(): return CliRunner() +@pytest.mark.parametrize( + "raw, expected", + [ + ("http://localhost:4000/", "http://localhost:4000"), + ("http://localhost:4000///", "http://localhost:4000"), + ("http://localhost:4000", "http://localhost:4000"), + ("https://host/gateway/", "https://host/gateway"), + ("", ""), + ], +) +def test_normalize_base_url(raw, expected): + """Trailing slashes are stripped while any path prefix is preserved""" + assert _normalize_base_url(raw) == expected + + +def test_cli_version_flag_normalizes_trailing_slash(cli_runner): + """--base-url with a trailing slash is normalised before the version output/health check""" + with patch( + "litellm.proxy.client.health.HealthManagementClient.get_server_version", + return_value="1.2.3", + ): + result = cli_runner.invoke(cli, ["--base-url", "http://localhost:4000/", "--version"]) + assert result.exit_code == 0 + assert "LiteLLM Proxy Server URL: http://localhost:4000" in result.output + assert "http://localhost:4000/" not in result.output + + +def test_cli_version_command_normalizes_trailing_slash(cli_runner): + """The `version` subcommand reports the normalised server URL from ctx.obj""" + with patch( + "litellm.proxy.client.health.HealthManagementClient.get_server_version", + return_value="1.2.3", + ): + result = cli_runner.invoke(cli, ["--base-url", "http://localhost:4000/", "version"]) + assert result.exit_code == 0 + assert "LiteLLM Proxy Server URL: http://localhost:4000" in result.output + assert "http://localhost:4000/" not in result.output + + def test_cli_version_flag(cli_runner): """Test that --version prints the correct version, server URL, and server version, and exits successfully""" with ( 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 045e15f8b8b..2c5aff27497 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py +++ b/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py @@ -37,6 +37,21 @@ from litellm.types.proxy.management_endpoints.ui_sso import ( ) +def _make_ui_sso_cache_mock() -> MagicMock: + """Build a cache mock whose async flow methods are awaitable. + + The CLI SSO flow persists/reads/deletes its session via the async cache + interface (async_get_cache / async_set_cache / async_delete_cache) so the + login works across proxy instances. Tests configure return values on these + the same way they would on a real async cache. + """ + cache = MagicMock() + cache.async_get_cache = AsyncMock(return_value=None) + cache.async_set_cache = AsyncMock() + cache.async_delete_cache = AsyncMock() + return cache + + def test_microsoft_sso_handler_openid_from_response_user_principal_name(): # Arrange # Create a mock response similar to what Microsoft SSO would return @@ -2151,7 +2166,7 @@ class TestCLIKeyRegenerationFlow: mock_request = MagicMock(spec=Request) mock_request.client = SimpleNamespace(host="127.0.0.1") mock_request.headers = {} - mock_cache = MagicMock() + mock_cache = _make_ui_sso_cache_mock() mock_cache.increment_cache.return_value = 1 with patch("litellm.proxy.proxy_server.user_api_key_cache", mock_cache): @@ -2163,8 +2178,8 @@ class TestCLIKeyRegenerationFlow: mock_cache.increment_cache.assert_called_once() assert mock_cache.increment_cache.call_args.kwargs["ttl"] == 60 - mock_cache.set_cache.assert_called_once() - flow_data = mock_cache.set_cache.call_args.kwargs["value"] + mock_cache.async_set_cache.assert_called_once() + flow_data = mock_cache.async_set_cache.call_args.kwargs["value"] assert flow_data["poll_secret_hash"] == _hash_cli_sso_secret( result["poll_secret"] ) @@ -2182,7 +2197,7 @@ class TestCLIKeyRegenerationFlow: mock_request = MagicMock(spec=Request) mock_request.client = SimpleNamespace(host="127.0.0.1") mock_request.headers = {} - mock_cache = MagicMock() + mock_cache = _make_ui_sso_cache_mock() mock_cache.increment_cache.return_value = 31 with patch("litellm.proxy.proxy_server.user_api_key_cache", mock_cache): @@ -2190,7 +2205,7 @@ class TestCLIKeyRegenerationFlow: await cli_sso_start(request=mock_request) assert exc_info.value.status_code == 429 - mock_cache.set_cache.assert_not_called() + mock_cache.async_set_cache.assert_not_called() @pytest.mark.asyncio async def test_cli_sso_start_returns_verification_uri_complete_when_enabled(self): @@ -2204,7 +2219,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 = _make_ui_sso_cache_mock() mock_cache.increment_cache.return_value = 1 with ( @@ -2238,7 +2253,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 = _make_ui_sso_cache_mock() mock_cache.increment_cache.return_value = 1 with ( @@ -2272,8 +2287,8 @@ class TestCLIKeyRegenerationFlow: mock_request = MagicMock(spec=Request) mock_request.base_url = "https://proxy.example.com/" - mock_cache = MagicMock() - mock_cache.get_cache.return_value = {"poll_secret_hash": "h"} + mock_cache = _make_ui_sso_cache_mock() + mock_cache.async_get_cache.return_value = {"poll_secret_hash": "h"} async def drive(enabled: bool): with ( @@ -2448,8 +2463,8 @@ class TestCLIKeyRegenerationFlow: ) mock_sso_result = {"user_email": "test@example.com", "user_id": "test-user-123"} - mock_cache = MagicMock() - mock_cache.get_cache.return_value = { + mock_cache = _make_ui_sso_cache_mock() + mock_cache.async_get_cache.return_value = { "poll_secret_hash": "poll-secret-hash", "user_code_hash": "user-code-hash", "sso_complete": False, @@ -2491,8 +2506,8 @@ class TestCLIKeyRegenerationFlow: mock_request.body = AsyncMock( return_value=b"user_code=ABCD-EFGH&browser_complete_token=browser-token" ) - mock_cache = MagicMock() - mock_cache.get_cache.return_value = { + mock_cache = _make_ui_sso_cache_mock() + mock_cache.async_get_cache.return_value = { "poll_secret_hash": _hash_cli_sso_secret("poll-secret"), "user_code_hash": _hash_cli_sso_secret( _normalize_cli_sso_user_code("ABCD-EFGH") @@ -2515,7 +2530,7 @@ class TestCLIKeyRegenerationFlow: ) assert result.status_code == 200 - flow_data = mock_cache.set_cache.call_args.kwargs["value"] + flow_data = mock_cache.async_set_cache.call_args.kwargs["value"] assert flow_data["user_code_verified"] is True @pytest.mark.asyncio @@ -2529,8 +2544,8 @@ class TestCLIKeyRegenerationFlow: mock_request = MagicMock(spec=Request) mock_request.body = AsyncMock(return_value=b"user_code=ABCD-EFGH") - mock_cache = MagicMock() - mock_cache.get_cache.return_value = { + mock_cache = _make_ui_sso_cache_mock() + mock_cache.async_get_cache.return_value = { "poll_secret_hash": _hash_cli_sso_secret("poll-secret"), "user_code_hash": _hash_cli_sso_secret( _normalize_cli_sso_user_code("ABCD-EFGH") @@ -2548,7 +2563,7 @@ class TestCLIKeyRegenerationFlow: ) assert exc_info.value.status_code == 400 - mock_cache.set_cache.assert_not_called() + mock_cache.async_set_cache.assert_not_called() @pytest.mark.asyncio async def test_cli_sso_complete_waits_for_callback_before_token_checks(self): @@ -2563,8 +2578,8 @@ class TestCLIKeyRegenerationFlow: mock_request.body = AsyncMock( return_value=b"user_code=ABCD-EFGH&browser_complete_token=browser-token" ) - mock_cache = MagicMock() - mock_cache.get_cache.return_value = { + mock_cache = _make_ui_sso_cache_mock() + mock_cache.async_get_cache.return_value = { "poll_secret_hash": _hash_cli_sso_secret("poll-secret"), "user_code_hash": _hash_cli_sso_secret( _normalize_cli_sso_user_code("ABCD-EFGH") @@ -2583,7 +2598,7 @@ class TestCLIKeyRegenerationFlow: assert exc_info.value.status_code == 400 assert exc_info.value.detail == "CLI login is not ready" mock_request.body.assert_not_awaited() - mock_cache.set_cache.assert_not_called() + mock_cache.async_set_cache.assert_not_called() @pytest.mark.asyncio async def test_cli_sso_callback_stores_session(self): @@ -2610,8 +2625,8 @@ class TestCLIKeyRegenerationFlow: mock_sso_result = {"user_email": "test@example.com", "user_id": "test-user-123"} # Mock cache - mock_cache = MagicMock() - mock_cache.get_cache.return_value = { + mock_cache = _make_ui_sso_cache_mock() + mock_cache.async_get_cache.return_value = { "poll_secret_hash": "poll-secret-hash", "user_code_hash": "user-code-hash", "sso_complete": False, @@ -2645,8 +2660,8 @@ class TestCLIKeyRegenerationFlow: ) # Assert - verify session was stored in cache - mock_cache.set_cache.assert_called_once() - call_args = mock_cache.set_cache.call_args + mock_cache.async_set_cache.assert_called_once() + call_args = mock_cache.async_set_cache.call_args # Verify cache key format assert "cli_sso_session:" in call_args.kwargs["key"] @@ -2692,8 +2707,8 @@ class TestCLIKeyRegenerationFlow: } # Mock cache - mock_cache = MagicMock() - mock_cache.get_cache.return_value = { + mock_cache = _make_ui_sso_cache_mock() + mock_cache.async_get_cache.return_value = { "poll_secret_hash": _hash_cli_sso_secret("poll-secret"), "sso_complete": True, "user_code_verified": True, @@ -2716,7 +2731,7 @@ class TestCLIKeyRegenerationFlow: assert "key" not in result # JWT should not be generated yet # Verify session was NOT deleted - mock_cache.delete_cache.assert_not_called() + mock_cache.async_delete_cache.assert_not_called() @pytest.mark.asyncio async def test_cli_poll_key_requires_poll_secret(self): @@ -2726,8 +2741,8 @@ class TestCLIKeyRegenerationFlow: cli_poll_key, ) - mock_cache = MagicMock() - mock_cache.get_cache.return_value = { + mock_cache = _make_ui_sso_cache_mock() + mock_cache.async_get_cache.return_value = { "poll_secret_hash": _hash_cli_sso_secret("poll-secret"), "sso_complete": True, "user_code_verified": True, @@ -2753,8 +2768,8 @@ class TestCLIKeyRegenerationFlow: cli_poll_key, ) - mock_cache = MagicMock() - mock_cache.get_cache.return_value = { + mock_cache = _make_ui_sso_cache_mock() + mock_cache.async_get_cache.return_value = { "poll_secret_hash": _hash_cli_sso_secret("poll-secret"), "sso_complete": True, "user_code_verified": False, @@ -2932,8 +2947,8 @@ class TestCLIKeyRegenerationFlow: ) # Mock cache - mock_cache = MagicMock() - mock_cache.get_cache.return_value = { + mock_cache = _make_ui_sso_cache_mock() + mock_cache.async_get_cache.return_value = { "poll_secret_hash": _hash_cli_sso_secret("poll-secret"), "sso_complete": True, "user_code_verified": True, @@ -2981,7 +2996,7 @@ class TestCLIKeyRegenerationFlow: assert jwt_call_args.kwargs["max_budget"] is None # Verify session was deleted after JWT generation - mock_cache.delete_cache.assert_called_once() + mock_cache.async_delete_cache.assert_called_once() @pytest.mark.asyncio async def test_cli_poll_key_does_not_cap_session_when_user_has_budget(self): @@ -3007,8 +3022,8 @@ class TestCLIKeyRegenerationFlow: models=["gpt-4"], max_budget=100.0, ) - mock_cache = MagicMock() - mock_cache.get_cache.return_value = { + mock_cache = _make_ui_sso_cache_mock() + mock_cache.async_get_cache.return_value = { "poll_secret_hash": _hash_cli_sso_secret("poll-secret"), "sso_complete": True, "user_code_verified": True, @@ -3069,8 +3084,8 @@ class TestCLIKeyRegenerationFlow: max_budget=None, ) mock_team = LiteLLM_TeamTableCachedObj(team_id="team-x", max_budget=None) - mock_cache = MagicMock() - mock_cache.get_cache.return_value = { + mock_cache = _make_ui_sso_cache_mock() + mock_cache.async_get_cache.return_value = { "poll_secret_hash": _hash_cli_sso_secret("poll-secret"), "sso_complete": True, "user_code_verified": True, @@ -3107,6 +3122,113 @@ class TestCLIKeyRegenerationFlow: ) +class TestCLISSOMultiInstanceCache: + """The CLI login is a multi-request flow whose intermediate session must be + visible to every proxy instance. These regressions pin the flow to the shared + Redis cache when configured (instead of the per-worker in-memory cache), which + is what breaks `lite login` when the gateway runs multiple instances.""" + + @pytest.mark.asyncio + async def test_cli_sso_start_persists_flow_to_shared_redis_cache(self): + from litellm.proxy.management_endpoints.ui_sso import ( + _get_cli_sso_flow_cache_key, + cli_sso_start, + ) + + mock_request = MagicMock(spec=Request) + mock_request.client = SimpleNamespace(host="127.0.0.1") + mock_request.headers = {} + + user_api_key_cache = _make_ui_sso_cache_mock() + user_api_key_cache.increment_cache.return_value = 1 + shared_redis_cache = _make_ui_sso_cache_mock() + + with ( + patch("litellm.proxy.proxy_server.user_api_key_cache", user_api_key_cache), + patch("litellm.proxy.proxy_server.redis_usage_cache", shared_redis_cache), + ): + result = await cli_sso_start(request=mock_request) + + shared_redis_cache.async_set_cache.assert_called_once() + user_api_key_cache.async_set_cache.assert_not_called() + assert shared_redis_cache.async_set_cache.call_args.kwargs[ + "key" + ] == _get_cli_sso_flow_cache_key(result["login_id"]) + + @pytest.mark.asyncio + async def test_cli_poll_reads_flow_from_shared_redis_cache(self): + from litellm.proxy.management_endpoints.ui_sso import ( + _get_cli_sso_flow_cache_key, + _hash_cli_sso_secret, + cli_poll_key, + ) + + poll_secret = "poll-secret" + session_key = "cli-shared-cache-session" + flow = { + "poll_secret_hash": _hash_cli_sso_secret(poll_secret), + "sso_complete": True, + "user_code_verified": True, + "session_data": { + "user_id": "cli-user", + "user_role": "internal_user", + "teams": [], + "models": ["gpt-4"], + }, + } + + user_api_key_cache = _make_ui_sso_cache_mock() + shared_redis_cache = _make_ui_sso_cache_mock() + shared_redis_cache.async_get_cache.return_value = flow + + with ( + patch("litellm.proxy.proxy_server.user_api_key_cache", user_api_key_cache), + patch("litellm.proxy.proxy_server.redis_usage_cache", shared_redis_cache), + patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), + patch( + "litellm.proxy.auth.auth_checks.get_user_object", + new=AsyncMock(side_effect=ValueError("User doesn't exist in db")), + ), + patch( + "litellm.proxy.auth.auth_checks.ExperimentalUIJWTToken.get_cli_jwt_auth_token", + return_value="cli-jwt-token", + ), + ): + result = await cli_poll_key( + key_id=session_key, x_litellm_cli_poll_secret=poll_secret + ) + + assert result["status"] == "ready" + assert result["key"] == "cli-jwt-token" + shared_redis_cache.async_get_cache.assert_called_with( + key=_get_cli_sso_flow_cache_key(session_key) + ) + user_api_key_cache.async_get_cache.assert_not_called() + shared_redis_cache.async_delete_cache.assert_called_once_with( + key=_get_cli_sso_flow_cache_key(session_key) + ) + user_api_key_cache.async_delete_cache.assert_not_called() + + @pytest.mark.asyncio + async def test_cli_sso_flow_falls_back_to_user_api_key_cache_without_redis(self): + from litellm.proxy.management_endpoints.ui_sso import cli_sso_start + + mock_request = MagicMock(spec=Request) + mock_request.client = SimpleNamespace(host="127.0.0.1") + mock_request.headers = {} + + user_api_key_cache = _make_ui_sso_cache_mock() + user_api_key_cache.increment_cache.return_value = 1 + + with ( + patch("litellm.proxy.proxy_server.user_api_key_cache", user_api_key_cache), + patch("litellm.proxy.proxy_server.redis_usage_cache", None), + ): + await cli_sso_start(request=mock_request) + + user_api_key_cache.async_set_cache.assert_called_once() + + class TestGetAppRolesFromIdToken: """Test the get_app_roles_from_id_token method""" @@ -4020,7 +4142,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 = _make_ui_sso_cache_mock() test_code_verifier = "test_code_verifier_abc123xyz" mock_cache.async_get_cache = AsyncMock( return_value={"code_verifier": test_code_verifier} @@ -4071,7 +4193,7 @@ class TestPKCEFunctionality: mock_sso.__exit__ = MagicMock(return_value=False) test_state = "test456" - mock_cache = MagicMock() + mock_cache = _make_ui_sso_cache_mock() mock_cache.async_set_cache = AsyncMock() @@ -4595,7 +4717,7 @@ class TestPKCEFunctionality: from litellm.proxy._types import ProxyException from litellm.proxy.management_endpoints.ui_sso import SSOAuthenticationHandler - mock_cache = MagicMock() + mock_cache = _make_ui_sso_cache_mock() mock_cache.async_get_cache = AsyncMock(return_value=None) # verifier not found mock_request = MagicMock(spec=Request) @@ -4721,7 +4843,7 @@ class TestPKCEFunctionality: from litellm.proxy.management_endpoints.ui_sso import SSOAuthenticationHandler # Cache returns an integer — unexpected format - mock_cache = MagicMock() + mock_cache = _make_ui_sso_cache_mock() mock_cache.async_get_cache = AsyncMock(return_value=12345) mock_cache.async_delete_cache = AsyncMock() @@ -4763,7 +4885,7 @@ class TestPKCEFunctionality: from litellm.proxy.management_endpoints.ui_sso import SSOAuthenticationHandler - mock_cache = MagicMock() + mock_cache = _make_ui_sso_cache_mock() mock_cache.async_get_cache = AsyncMock(return_value=None) # verifier not found mock_request = MagicMock(spec=Request) @@ -4851,7 +4973,7 @@ class TestPKCEFunctionality: from litellm.proxy.management_endpoints.ui_sso import SSOAuthenticationHandler # Cache returns an integer — unexpected format - mock_cache = MagicMock() + mock_cache = _make_ui_sso_cache_mock() mock_cache.async_get_cache = AsyncMock(return_value=12345) mock_cache.async_delete_cache = AsyncMock() @@ -4903,7 +5025,7 @@ class TestPKCEFunctionality: from litellm.proxy.management_endpoints.ui_sso import SSOAuthenticationHandler legacy_verifier = "legacy_plain_string_verifier_abc123" - mock_cache = MagicMock() + mock_cache = _make_ui_sso_cache_mock() mock_cache.async_get_cache = AsyncMock(return_value=legacy_verifier) mock_request = MagicMock(spec=Request) @@ -6187,8 +6309,8 @@ class TestCliSsoAttributionMetadata: provider="generic", team_ids=[], ) - mock_cache = MagicMock() - mock_cache.get_cache.return_value = { + mock_cache = _make_ui_sso_cache_mock() + mock_cache.async_get_cache.return_value = { "poll_secret_hash": "poll-secret-hash", "user_code_hash": "user-code-hash", "sso_complete": False, @@ -6228,8 +6350,8 @@ class TestCliSsoAttributionMetadata: mock_request = MagicMock(spec=Request) mock_request.base_url = "http://internal-proxy.local/" - mock_cache = MagicMock() - mock_cache.get_cache.return_value = { + mock_cache = _make_ui_sso_cache_mock() + mock_cache.async_get_cache.return_value = { "poll_secret_hash": "poll-secret-hash", "user_code_hash": "user-code-hash", "sso_complete": False, @@ -6297,8 +6419,8 @@ class TestCliSsoAttributionMetadata: "user_id": "test-user-123", "employment_type": "contractor", } - mock_cache = MagicMock() - mock_cache.get_cache.return_value = { + mock_cache = _make_ui_sso_cache_mock() + mock_cache.async_get_cache.return_value = { "poll_secret_hash": "poll-secret-hash", "user_code_hash": "user-code-hash", "sso_complete": False, @@ -6337,7 +6459,7 @@ class TestCliSsoAttributionMetadata: result=mock_sso_result, ) - flow_data = mock_cache.set_cache.call_args.kwargs["value"] + flow_data = mock_cache.async_set_cache.call_args.kwargs["value"] assert flow_data["session_data"]["attribution_metadata"] == { "acme_employment_type": "contractor" } @@ -6366,8 +6488,8 @@ class TestCliSsoAttributionMetadata: "org": {"cost_center": "CC-42"}, }, } - mock_cache = MagicMock() - mock_cache.get_cache.return_value = { + mock_cache = _make_ui_sso_cache_mock() + mock_cache.async_get_cache.return_value = { "poll_secret_hash": _hash_cli_sso_secret("poll-secret"), "sso_complete": True, "user_code_verified": True, @@ -7225,8 +7347,8 @@ async def test_cli_poll_key_tolerates_missing_user_row(): "models": ["gpt-4"], } - mock_cache = MagicMock() - mock_cache.get_cache.return_value = { + mock_cache = _make_ui_sso_cache_mock() + mock_cache.async_get_cache.return_value = { "poll_secret_hash": _hash_cli_sso_secret("poll-secret"), "sso_complete": True, "user_code_verified": True,