mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-14 23:21:35 +00:00
fix(cli): normalize trailing-slash base url and share CLI SSO login flow across proxy instances
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
mateo
parent
b340a2693c
commit
c6be1a1fd1
5 changed files with 306 additions and 69 deletions
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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 = {
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -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 (
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue