tiennm99/litellm
1
0
Fork 0
mirror of https://github.com/tiennm99/litellm.git synced 2026-07-08 03:14:09 +00:00
Code Issues Packages Projects Releases Wiki Activity

tighten cli sso session flow

Browse Source
This commit is contained in:
user
2026-04-29 17:13:25 -07:00
parent 9bc317b4d0
commit 88d8a80761
6 changed files with 605 additions and 167 deletions
Download Patch File Download Diff File Expand all files Collapse all files
litellm/constants.py
+1
View File
@@ -1421,6 +1421,7 @@ LITELLM_PROXY_ADMIN_NAME = "default_user_id"
LITELLM_CLI_SOURCE_IDENTIFIER = "litellm-cli"
LITELLM_CLI_SESSION_TOKEN_PREFIX = "litellm-session-token"
CLI_SSO_SESSION_CACHE_KEY_PREFIX = "cli_sso_session"
CLI_SSO_SESSION_TTL_SECONDS = 600
CLI_JWT_TOKEN_NAME = "cli-jwt-token"
# Support both CLI_JWT_EXPIRATION_HOURS and LITELLM_CLI_JWT_EXPIRATION_HOURS for backwards compatibility
CLI_JWT_EXPIRATION_HOURS = int(
litellm/proxy/client/README.md
+17 -16
View File
@@ -313,23 +313,24 @@ sequenceDiagram
participant Proxy as LiteLLM Proxy
participant SSO as SSO Provider
CLI->>CLI: Generate key ID (sk-uuid)
CLI->>Browser: Open /sso/key/generate?source=litellm-cli&key=sk-uuid
CLI->>Proxy: POST /sso/cli/start
Proxy->>CLI: Return login_id, poll_secret, user_code
CLI->>Browser: Open /sso/key/generate?source=litellm-cli&key=login_id
Browser->>Proxy: GET /sso/key/generate?source=litellm-cli&key=sk-uuid
Proxy->>Proxy: Set cli_state = litellm-session-token:sk-uuid
Proxy->>SSO: Redirect with state=litellm-session-token:sk-uuid
Browser->>Proxy: GET /sso/key/generate?source=litellm-cli&key=login_id
Proxy->>Proxy: Set cli_state = litellm-session-token:login_id
Proxy->>SSO: Redirect with state=litellm-session-token:login_id
SSO->>Browser: Show login page
Browser->>SSO: User authenticates
SSO->>Proxy: Redirect to /sso/callback?state=litellm-session-token:sk-uuid
SSO->>Proxy: Redirect to /sso/callback?state=litellm-session-token:login_id
Proxy->>Proxy: Check if state starts with "litellm-session-token:"
Proxy->>Proxy: Generate API key with ID=sk-uuid
Proxy->>Browser: Show success page
Proxy->>Browser: Prompt for user_code
Browser->>Proxy: POST /sso/cli/complete/login_id
CLI->>Proxy: Poll /sso/cli/poll/sk-uuid
Proxy->>CLI: Return {"status": "ready", "key": "sk-uuid"}
CLI->>Proxy: Poll /sso/cli/poll/login_id with poll_secret header
Proxy->>CLI: Return {"status": "ready", "key": "jwt"}
CLI->>CLI: Save key to ~/.litellm/token.json
```
@@ -343,13 +344,13 @@ The CLI provides three authentication commands:
### Authentication Flow Steps
1. **Generate Session ID**: CLI generates a unique key ID (`sk-{uuid}`)
2. **Open Browser**: CLI opens browser to `/sso/key/generate` with CLI source and key parameters
3. **SSO Redirect**: Proxy sets the formatted state (`litellm-session-token:sk-uuid`) as OAuth state parameter and redirects to SSO provider
1. **Start Session**: CLI creates a short-lived login session with `/sso/cli/start`
2. **Open Browser**: CLI opens browser to `/sso/key/generate` with CLI source and login ID parameters
3. **SSO Redirect**: Proxy sets the formatted state (`litellm-session-token:{login_id}`) as OAuth state parameter and redirects to SSO provider
4. **User Authentication**: User completes SSO authentication in browser
5. **Callback Processing**: SSO provider redirects back to proxy with state parameter
6. **Key Generation**: Proxy detects CLI login (state starts with "litellm-session-token:") and generates API key with pre-specified ID
7. **Polling**: CLI polls `/sso/cli/poll/{key_id}` endpoint until key is ready
6. **User Code Verification**: Browser confirms the verification code shown in the CLI
7. **Polling**: CLI polls `/sso/cli/poll/{login_id}` with the polling secret header until the JWT is ready
8. **Token Storage**: CLI saves the authentication token to `~/.litellm/token.json`
### Benefits of This Approach
@@ -357,7 +358,7 @@ The CLI provides three authentication commands:
- **No Local Server**: No need to run a local callback server
- **Standard OAuth**: Uses OAuth 2.0 state parameter correctly
- **Remote Compatible**: Works with remote proxy servers
- **Secure**: Uses UUID session identifiers
- **Secure**: Keeps the polling secret out of the browser handoff
- **Simple Setup**: No additional OAuth redirect URL configuration needed
### Token Storage
litellm/proxy/client/cli/commands/auth.py
+36 -18
View File
@@ -5,6 +5,7 @@ import time
import webbrowser
from pathlib import Path
from typing import Any, Dict, List, Optional
from urllib.parse import urlencode
import click
import requests
@@ -241,7 +242,7 @@ def prompt_team_selection(teams: List[Dict[str, Any]]) -> Optional[Dict[str, Any
def prompt_team_selection_fallback(
teams: List[Dict[str, Any]]
teams: List[Dict[str, Any]],
) -> Optional[Dict[str, Any]]:
"""Fallback team selection for non-interactive environments"""
if not teams:
@@ -279,6 +280,7 @@ def prompt_team_selection_fallback(
def _poll_for_ready_data(
url: str,
*,
headers: Optional[Dict[str, str]] = None,
total_timeout: int = 300,
poll_interval: int = 2,
request_timeout: int = 10,
@@ -291,7 +293,7 @@ def _poll_for_ready_data(
) -> Optional[Dict[str, Any]]:
for attempt in range(total_timeout // poll_interval):
try:
response = requests.get(url, timeout=request_timeout)
response = requests.get(url, headers=headers, timeout=request_timeout)
if response.status_code == 200:
data = response.json()
status = data.get("status")
@@ -346,7 +348,23 @@ def _normalize_teams(teams, team_details):
return []
def _poll_for_authentication(base_url: str, key_id: str) -> Optional[dict]:
def _start_cli_sso_flow(base_url: str) -> Dict[str, Any]:
response = requests.post(f"{base_url}/sso/cli/start", timeout=10)
response.raise_for_status()
data = response.json()
required_fields = ("login_id", "poll_secret", "user_code")
if not all(isinstance(data.get(field), str) for field in required_fields):
raise ValueError("Invalid CLI SSO start response")
return data
def _get_cli_sso_poll_headers(poll_secret: str) -> Dict[str, str]:
return {"x-litellm-cli-poll-secret": poll_secret}
def _poll_for_authentication(
base_url: str, key_id: str, poll_secret: str
) -> Optional[dict]:
"""
Poll the server for authentication completion and handle team selection.
@@ -356,6 +374,7 @@ def _poll_for_authentication(base_url: str, key_id: str) -> Optional[dict]:
poll_url = f"{base_url}/sso/cli/poll/{key_id}"
data = _poll_for_ready_data(
poll_url,
headers=_get_cli_sso_poll_headers(poll_secret),
pending_message="Still waiting for authentication...",
)
if not data:
@@ -373,6 +392,7 @@ def _poll_for_authentication(base_url: str, key_id: str) -> Optional[dict]:
jwt_with_team = _handle_team_selection_during_polling(
base_url=base_url,
key_id=key_id,
poll_secret=poll_secret,
teams=normalized_teams,
)
@@ -410,7 +430,7 @@ def _poll_for_authentication(base_url: str, key_id: str) -> Optional[dict]:
def _handle_team_selection_during_polling(
base_url: str, key_id: str, teams: List[Dict[str, Any]]
base_url: str, key_id: str, poll_secret: str, teams: List[Dict[str, Any]]
) -> Optional[str]:
"""
Handle team selection and re-poll with selected team_id.
@@ -441,6 +461,7 @@ def _handle_team_selection_during_polling(
poll_url = f"{base_url}/sso/cli/poll/{key_id}?team_id={team_id}"
data = _poll_for_ready_data(
poll_url,
headers=_get_cli_sso_poll_headers(poll_secret),
pending_message="Still waiting for team authentication...",
other_status_message="Waiting for team authentication to complete...",
http_error_log_every=10,
@@ -514,29 +535,24 @@ def _render_and_prompt_for_team_selection(teams: List[Dict[str, Any]]) -> Option
@click.pass_context
def login(ctx: click.Context):
"""Login to LiteLLM proxy using SSO authentication"""
from litellm._uuid import uuid
from litellm.constants import LITELLM_CLI_SOURCE_IDENTIFIER
from litellm.proxy.client.cli.interface import show_commands
base_url = ctx.obj["base_url"]
# Check if we have an existing key to regenerate
existing_key = get_stored_api_key()
# Generate unique key ID for this login session
key_id = f"sk-{str(uuid.uuid4())}"
try:
# Construct SSO login URL with CLI source and pre-generated key
sso_url = f"{base_url}/sso/key/generate?source={LITELLM_CLI_SOURCE_IDENTIFIER}&key={key_id}"
cli_sso_flow = _start_cli_sso_flow(base_url=base_url)
key_id = cli_sso_flow["login_id"]
poll_secret = cli_sso_flow["poll_secret"]
user_code = cli_sso_flow["user_code"]
# If we have an existing key, include it as a parameter to the login endpoint
# The server will encode it in the OAuth state parameter for the SSO flow
if existing_key:
sso_url += f"&existing_key={existing_key}"
sso_url = f"{base_url}/sso/key/generate?" + urlencode(
{"source": LITELLM_CLI_SOURCE_IDENTIFIER, "key": key_id}
)
click.echo(f"Opening browser to: {sso_url}")
click.echo("Please complete the SSO authentication in your browser...")
click.echo(f"Verification code: {user_code}")
click.echo(f"Session ID: {key_id}")
# Open browser
@@ -545,7 +561,9 @@ def login(ctx: click.Context):
# Poll for authentication completion
click.echo("Waiting for authentication...")
auth_result = _poll_for_authentication(base_url=base_url, key_id=key_id)
auth_result = _poll_for_authentication(
base_url=base_url, key_id=key_id, poll_secret=poll_secret
)
if auth_result:
api_key = auth_result["api_key"]
litellm/proxy/management_endpoints/ui_sso.py
+251 -52
View File
@@ -14,6 +14,7 @@ import hashlib
import inspect
import os
import secrets
from html import escape
from copy import deepcopy
from typing import (
TYPE_CHECKING,
@@ -27,13 +28,13 @@ from typing import (
Union,
cast,
)
from urllib.parse import urlencode, urlparse
from urllib.parse import parse_qs, urlencode, urlparse
if TYPE_CHECKING:
import httpx
import jwt
from fastapi import APIRouter, Depends, HTTPException, Request, status
from fastapi import APIRouter, Depends, Header, HTTPException, Request, status
from fastapi.responses import RedirectResponse
import litellm
@@ -41,6 +42,9 @@ from litellm._logging import verbose_proxy_logger
from litellm._uuid import uuid
from litellm.caching import DualCache
from litellm.constants import (
CLI_SSO_SESSION_CACHE_KEY_PREFIX,
CLI_SSO_SESSION_TTL_SECONDS,
LITELLM_CLI_SOURCE_IDENTIFIER,
LITELLM_UI_SESSION_DURATION,
MAX_SPENDLOG_ROWS_TO_QUERY,
MICROSOFT_USER_DISPLAY_NAME_ATTRIBUTE,
@@ -123,6 +127,207 @@ router = APIRouter()
# Metadata fields (token_type, expires_in, scope) are intentionally kept so
# response convertors see the same fields in the PKCE path as in the non-PKCE path.
_OAUTH_TOKEN_FIELDS = frozenset({"access_token", "id_token", "refresh_token"})
_CLI_SSO_FLOW_CACHE_KEY_PREFIX = f"{CLI_SSO_SESSION_CACHE_KEY_PREFIX}:flow"
_CLI_SSO_USER_CODE_ALPHABET = "ABCDEFGHJKLMNPQRSTUVWXYZ23456789"
def _hash_cli_sso_secret(secret: str) -> str:
return hashlib.sha256(secret.encode("utf-8")).hexdigest()
def _normalize_cli_sso_user_code(user_code: str) -> str:
return "".join(ch for ch in user_code.upper() if ch.isalnum())
def _generate_cli_sso_user_code() -> str:
user_code = "".join(secrets.choice(_CLI_SSO_USER_CODE_ALPHABET) for _ in range(8))
return f"{user_code[:4]}-{user_code[4:]}"
def _get_cli_sso_flow_cache_key(login_id: str) -> str:
return f"{_CLI_SSO_FLOW_CACHE_KEY_PREFIX}:{login_id}"
def _is_valid_cli_sso_login_id(login_id: Optional[str]) -> bool:
return (
isinstance(login_id, str)
and login_id.startswith("cli-")
and 16 <= len(login_id) <= 128
)
def _get_cli_sso_flow_or_raise(login_id: Optional[str], cache: DualCache) -> 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)
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(
key=_get_cli_sso_flow_cache_key(login_id),
value=flow,
ttl=CLI_SSO_SESSION_TTL_SECONDS,
)
def _verify_cli_sso_poll_secret(flow: dict, poll_secret: Optional[str]) -> bool:
expected_poll_secret_hash = flow.get("poll_secret_hash")
if not isinstance(expected_poll_secret_hash, str) or not isinstance(
poll_secret, str
):
return False
supplied_poll_secret_hash = _hash_cli_sso_secret(poll_secret)
return secrets.compare_digest(supplied_poll_secret_hash, expected_poll_secret_hash)
def _render_cli_sso_verification_page(
verify_url: str, browser_complete_token: str
) -> str:
escaped_verify_url = escape(verify_url, quote=True)
escaped_browser_complete_token = escape(browser_complete_token, quote=True)
return f"""
<!doctype html>
<html>
<head>
<title>LiteLLM CLI Login</title>
<style>
body {{
font-family: -apple-system, BlinkMacSystemFont, "Segoe UI", sans-serif;
margin: 0;
min-height: 100vh;
display: flex;
align-items: center;
justify-content: center;
background: #f8fafc;
color: #0f172a;
}}
main {{
width: min(420px, calc(100vw - 32px));
background: #ffffff;
border: 1px solid #e2e8f0;
border-radius: 8px;
padding: 28px;
box-shadow: 0 12px 32px rgba(15, 23, 42, 0.08);
}}
h1 {{ font-size: 22px; margin: 0 0 12px; }}
p {{ line-height: 1.5; margin: 0 0 18px; color: #334155; }}
label {{ display: block; font-weight: 600; margin-bottom: 8px; }}
input {{
box-sizing: border-box;
width: 100%;
padding: 12px;
border: 1px solid #cbd5e1;
border-radius: 6px;
font-size: 20px;
letter-spacing: 0.08em;
text-transform: uppercase;
}}
button {{
margin-top: 16px;
width: 100%;
padding: 12px;
border: 0;
border-radius: 6px;
background: #0f172a;
color: #ffffff;
font-weight: 600;
cursor: pointer;
}}
</style>
</head>
<body>
<main>
<h1>Complete CLI Login</h1>
<p>Enter the verification code shown in your terminal to finish this login.</p>
<form method="post" action="{escaped_verify_url}">
<input type="hidden" name="browser_complete_token" value="{escaped_browser_complete_token}" />
<label for="user_code">Verification code</label>
<input id="user_code" name="user_code" autocomplete="one-time-code" required autofocus />
<button type="submit">Continue</button>
</form>
</main>
</body>
</html>
"""
@router.post("/sso/cli/start", tags=["experimental"], include_in_schema=False)
async def cli_sso_start():
from litellm.proxy.proxy_server import user_api_key_cache
login_id = f"cli-{secrets.token_urlsafe(24)}"
poll_secret = secrets.token_urlsafe(32)
user_code = _generate_cli_sso_user_code()
flow = {
"poll_secret_hash": _hash_cli_sso_secret(poll_secret),
"user_code_hash": _hash_cli_sso_secret(_normalize_cli_sso_user_code(user_code)),
"sso_complete": False,
"user_code_verified": False,
"session_data": None,
}
_set_cli_sso_flow(login_id=login_id, cache=user_api_key_cache, flow=flow)
return {
"login_id": login_id,
"poll_secret": poll_secret,
"user_code": user_code,
"expires_in": CLI_SSO_SESSION_TTL_SECONDS,
}
@router.post(
"/sso/cli/complete/{login_id}", tags=["experimental"], include_in_schema=False
)
async def cli_sso_complete(request: Request, login_id: str):
from fastapi.responses import HTMLResponse
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
flow = _get_cli_sso_flow_or_raise(login_id=login_id, cache=user_api_key_cache)
body = (await request.body()).decode("utf-8")
form_values = parse_qs(body)
supplied_user_code = (form_values.get("user_code") or [""])[0]
supplied_browser_complete_token = (
form_values.get("browser_complete_token") or [""]
)[0]
supplied_user_code_hash = _hash_cli_sso_secret(
_normalize_cli_sso_user_code(supplied_user_code)
)
supplied_browser_complete_token_hash = _hash_cli_sso_secret(
supplied_browser_complete_token
)
expected_user_code_hash = flow.get("user_code_hash")
if not isinstance(expected_user_code_hash, str) or not secrets.compare_digest(
supplied_user_code_hash, expected_user_code_hash
):
raise HTTPException(status_code=400, detail="Invalid verification code")
expected_browser_complete_token_hash = flow.get("browser_complete_token_hash")
if not isinstance(
expected_browser_complete_token_hash, str
) or not secrets.compare_digest(
supplied_browser_complete_token_hash, expected_browser_complete_token_hash
):
raise HTTPException(status_code=400, detail="Invalid verification code")
if not flow.get("sso_complete") or not flow.get("session_data"):
raise HTTPException(status_code=400, detail="CLI login is not ready")
flow["user_code_verified"] = True
_set_cli_sso_flow(login_id=login_id, cache=user_api_key_cache, flow=flow)
html_content = render_cli_sso_success_page()
return HTMLResponse(content=html_content, status_code=200)
def normalize_email(email: Optional[str]) -> Optional[str]:
@@ -333,6 +538,7 @@ async def google_login(
from litellm.proxy.proxy_server import (
premium_user,
prisma_client,
user_api_key_cache,
user_custom_ui_sso_sign_in_handler,
)
@@ -382,14 +588,15 @@ async def google_login(
redirect_url = SSOAuthenticationHandler.get_redirect_url_for_sso(
request=request,
sso_callback_route="sso/callback",
existing_key=existing_key,
)
# Store CLI key in state for OAuth flow
if source == LITELLM_CLI_SOURCE_IDENTIFIER:
_get_cli_sso_flow_or_raise(login_id=key, cache=user_api_key_cache)
# Store CLI login handle in state for OAuth flow
cli_state: Optional[str] = SSOAuthenticationHandler._get_cli_state(
source=source,
key=key,
existing_key=existing_key,
)
# check if user defined a custom auth sso sign in handler, if yes, use it
@@ -1392,18 +1599,12 @@ async def auth_callback(request: Request, state: Optional[str] = None): # noqa:
)
if state and state.startswith(f"{LITELLM_CLI_SESSION_TOKEN_PREFIX}:"):
# Extract the key ID and existing_key from the state
# State format: {PREFIX}:{key}:{existing_key} or {PREFIX}:{key}
state_parts = state.split(":", 2) # Split into max 3 parts
# State format: {PREFIX}:{login_id}
state_parts = state.split(":", 1)
key_id = state_parts[1] if len(state_parts) > 1 else None
existing_key = state_parts[2] if len(state_parts) > 2 else None
verbose_proxy_logger.info(
f"CLI SSO callback detected for key: {key_id}, existing_key: {existing_key}"
)
return await cli_sso_callback(
request=request, key=key_id, existing_key=existing_key, result=result
)
verbose_proxy_logger.info("CLI SSO callback detected")
return await cli_sso_callback(request=request, key=key_id, result=result)
# Control-plane cross-origin: read return_to from cookie.
# Starlette's cookie_parser already handles RFC 2109 unquoting.
@@ -1424,13 +1625,10 @@ async def auth_callback(request: Request, state: Optional[str] = None): # noqa:
async def cli_sso_callback(
request: Request,
key: Optional[str] = None,
existing_key: Optional[str] = None,
result: Optional[Union[OpenID, dict]] = None,
):
"""CLI SSO callback - stores session info for JWT generation on polling"""
verbose_proxy_logger.info(
f"CLI SSO callback for key: {key}, existing_key: {existing_key}"
)
verbose_proxy_logger.info("CLI SSO callback")
from litellm.proxy.proxy_server import (
prisma_client,
@@ -1438,11 +1636,7 @@ async def cli_sso_callback(
user_api_key_cache,
)
if not key or not key.startswith("sk-"):
raise HTTPException(
status_code=400,
detail="Invalid key parameter. Must be a valid key ID starting with 'sk-'",
)
flow = _get_cli_sso_flow_or_raise(login_id=key, cache=user_api_key_cache)
if prisma_client is None:
raise HTTPException(
@@ -1480,9 +1674,6 @@ async def cli_sso_callback(
status_code=500, detail="Failed to retrieve user information from SSO"
)
# Store session info in cache (10 min TTL)
from litellm.constants import CLI_SSO_SESSION_CACHE_KEY_PREFIX
# Get all teams from user_info - CLI will let user select which one
teams: List[str] = []
if hasattr(user_info, "teams") and user_info.teams:
@@ -1523,21 +1714,25 @@ async def cli_sso_callback(
"team_details": team_details,
}
cache_key = f"{CLI_SSO_SESSION_CACHE_KEY_PREFIX}:{key}"
user_api_key_cache.set_cache(key=cache_key, value=session_data, ttl=600)
flow["session_data"] = session_data
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=cast(str, key), 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)}"
)
# Return success page
from fastapi.responses import HTMLResponse
from litellm.proxy.common_utils.html_forms.cli_sso_success import (
render_cli_sso_success_page,
verify_url = str(request.url_for("cli_sso_complete", login_id=key))
html_content = _render_cli_sso_verification_page(
verify_url=verify_url,
browser_complete_token=browser_complete_token,
)
html_content = render_cli_sso_success_page()
return HTMLResponse(content=html_content, status_code=200)
except Exception as e:
@@ -1548,7 +1743,11 @@ async def cli_sso_callback(
@router.get("/sso/cli/poll/{key_id}", tags=["experimental"], include_in_schema=False)
async def cli_poll_key(key_id: str, team_id: Optional[str] = None):
async def cli_poll_key(
key_id: str,
team_id: Optional[str] = None,
x_litellm_cli_poll_secret: Optional[str] = Header(default=None),
):
"""
CLI polling endpoint - retrieves session from cache and generates JWT.
@@ -1557,22 +1756,25 @@ async def cli_poll_key(key_id: str, team_id: Optional[str] = None):
2. Second poll (with team_id): Generates JWT with selected team and deletes session
Args:
key_id: The session key ID
key_id: The CLI login session ID
team_id: Optional team ID to assign to the JWT. If provided, must be one of user's teams.
"""
from litellm.constants import CLI_SSO_SESSION_CACHE_KEY_PREFIX
from litellm.proxy.auth.auth_checks import ExperimentalUIJWTToken
from litellm.proxy.proxy_server import user_api_key_cache
if not key_id.startswith("sk-"):
raise HTTPException(status_code=400, detail="Invalid key ID format")
try:
# Look up session in cache
cache_key = f"{CLI_SSO_SESSION_CACHE_KEY_PREFIX}:{key_id}"
session_data = user_api_key_cache.get_cache(key=cache_key)
flow = _get_cli_sso_flow_or_raise(login_id=key_id, cache=user_api_key_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")
if session_data:
if not flow.get("sso_complete") or not flow.get("user_code_verified"):
return {"status": "pending"}
session_data = flow.get("session_data")
if isinstance(session_data, dict):
user_teams = session_data.get("teams", [])
user_team_details = session_data.get("team_details")
user_id = session_data["user_id"]
@@ -1632,7 +1834,7 @@ async def cli_poll_key(key_id: str, team_id: Optional[str] = None):
)
# Delete cache entry (single-use)
user_api_key_cache.delete_cache(key=cache_key)
user_api_key_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}"
@@ -1650,6 +1852,8 @@ async def cli_poll_key(key_id: str, team_id: Optional[str] = None):
else:
return {"status": "pending"}
except HTTPException:
raise
except Exception as e:
verbose_proxy_logger.error(f"Error polling for CLI JWT: {e}")
raise HTTPException(
@@ -2393,20 +2597,15 @@ class SSOAuthenticationHandler:
This is used to authenticate through the CLI login flow.
The state parameter format is: {PREFIX}:{key}:{existing_key}
- If existing_key is provided, it's included in the state
The state parameter format is: {PREFIX}:{login_id}
- The state parameter is used to pass data through the OAuth flow without changing the callback URL
"""
from litellm.constants import (
LITELLM_CLI_SESSION_TOKEN_PREFIX,
LITELLM_CLI_SOURCE_IDENTIFIER,
)
if source == LITELLM_CLI_SOURCE_IDENTIFIER and key:
if existing_key:
return f"{LITELLM_CLI_SESSION_TOKEN_PREFIX}:{key}:{existing_key}"
else:
return f"{LITELLM_CLI_SESSION_TOKEN_PREFIX}:{key}"
return f"{LITELLM_CLI_SESSION_TOKEN_PREFIX}:{key}"
else:
return None
tests/test_litellm/proxy/client/cli/test_auth_commands.py
+52 -16
View File
@@ -1,17 +1,15 @@
import json
import os
import sys
import tempfile
import time
from pathlib import Path
from unittest.mock import MagicMock, Mock, mock_open, patch
from unittest.mock import Mock, mock_open, patch
sys.path.insert(
0, os.path.abspath("../../..")
) # Adds the parent directory to the system path
import pytest
from click.testing import CliRunner
from litellm.proxy.client.cli.commands.auth import (
@@ -26,6 +24,22 @@ from litellm.proxy.client.cli.commands.auth import (
)
def _mock_cli_sso_start_response(
login_id: str = "cli-session-uuid-456",
poll_secret: str = "poll-secret",
user_code: str = "ABCD-EFGH",
) -> Mock:
mock_response = Mock()
mock_response.status_code = 200
mock_response.json.return_value = {
"login_id": login_id,
"poll_secret": poll_secret,
"user_code": user_code,
}
mock_response.raise_for_status = Mock()
return mock_response
class TestTokenUtilities:
"""Test token file utility functions"""
@@ -243,12 +257,15 @@ class TestLoginCommand:
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.interface.show_commands"
) as mock_show_commands,
patch("litellm._uuid.uuid.uuid4", return_value="test-uuid-123"),
):
result = self.runner.invoke(login, obj=mock_context.obj)
@@ -261,7 +278,13 @@ class TestLoginCommand:
mock_browser.assert_called_once()
call_args = mock_browser.call_args[0][0]
assert "https://test.example.com/sso/key/generate" in call_args
assert "sk-test-uuid-123" in call_args
assert "cli-test-uuid-123" in call_args
assert "Verification code: ABCD-EFGH" in result.output
mock_post.assert_called_once()
mock_get.assert_called()
assert mock_get.call_args.kwargs["headers"] == {
"x-litellm-cli-poll-secret": "poll-secret"
}
# Verify JWT was saved
mock_save.assert_called_once()
@@ -284,9 +307,9 @@ class TestLoginCommand:
with (
patch("webbrowser.open"),
patch("requests.post", return_value=_mock_cli_sso_start_response()),
patch("requests.get", return_value=mock_response),
patch("time.sleep") as mock_sleep,
patch("litellm._uuid.uuid.uuid4", return_value="test-uuid-123"),
patch("time.sleep"),
):
# Mock time.sleep to avoid actual delays in tests
@@ -306,9 +329,9 @@ class TestLoginCommand:
with (
patch("webbrowser.open"),
patch("requests.post", return_value=_mock_cli_sso_start_response()),
patch("requests.get", return_value=mock_response),
patch("time.sleep"),
patch("litellm._uuid.uuid.uuid4", return_value="test-uuid-123"),
):
result = self.runner.invoke(login, obj=mock_context.obj)
@@ -325,12 +348,12 @@ class TestLoginCommand:
with (
patch("webbrowser.open"),
patch("requests.post", return_value=_mock_cli_sso_start_response()),
patch(
"requests.get",
side_effect=requests.RequestException("Connection failed"),
),
patch("time.sleep"),
patch("litellm._uuid.uuid.uuid4", return_value="test-uuid-123"),
):
result = self.runner.invoke(login, obj=mock_context.obj)
@@ -345,8 +368,8 @@ class TestLoginCommand:
with (
patch("webbrowser.open"),
patch("requests.post", return_value=_mock_cli_sso_start_response()),
patch("requests.get", side_effect=KeyboardInterrupt),
patch("litellm._uuid.uuid.uuid4", return_value="test-uuid-123"),
):
result = self.runner.invoke(login, obj=mock_context.obj)
@@ -369,9 +392,9 @@ class TestLoginCommand:
with (
patch("webbrowser.open"),
patch("requests.post", return_value=_mock_cli_sso_start_response()),
patch("requests.get", return_value=mock_response),
patch("time.sleep"),
patch("litellm._uuid.uuid.uuid4", return_value="test-uuid-123"),
):
result = self.runner.invoke(login, obj=mock_context.obj)
@@ -386,8 +409,8 @@ class TestLoginCommand:
with (
patch("webbrowser.open"),
patch("requests.post", return_value=_mock_cli_sso_start_response()),
patch("requests.get", side_effect=ValueError("Invalid value")),
patch("litellm._uuid.uuid.uuid4", return_value="test-uuid-123"),
):
result = self.runner.invoke(login, obj=mock_context.obj)
@@ -556,6 +579,12 @@ class TestCLIKeyRegenerationFlow:
# Simulate user selecting team #2 (team-beta)
with (
patch("webbrowser.open") as mock_browser,
patch(
"requests.post",
return_value=_mock_cli_sso_start_response(
login_id="cli-session-uuid-456"
),
),
patch(
"requests.get", side_effect=[mock_first_response, mock_second_response]
) as mock_get,
@@ -563,7 +592,6 @@ class TestCLIKeyRegenerationFlow:
patch(
"litellm.proxy.client.cli.interface.show_commands"
) as mock_show_commands,
patch("litellm._uuid.uuid.uuid4", return_value="session-uuid-456"),
patch("click.prompt", return_value="2"),
): # User selects index 2
@@ -585,8 +613,11 @@ class TestCLIKeyRegenerationFlow:
# First poll should be without team_id
first_poll_url = mock_get.call_args_list[0][0][0]
assert "sk-session-uuid-456" in first_poll_url
assert "cli-session-uuid-456" in first_poll_url
assert "team_id=" not in first_poll_url
assert mock_get.call_args_list[0].kwargs["headers"] == {
"x-litellm-cli-poll-secret": "poll-secret"
}
# Second poll should include team_id=team-beta
second_poll_url = mock_get.call_args_list[1][0][0]
@@ -621,10 +652,15 @@ class TestCLIKeyRegenerationFlow:
with (
patch("webbrowser.open") as mock_browser,
patch(
"requests.post",
return_value=_mock_cli_sso_start_response(
login_id="cli-session-uuid-solo"
),
),
patch("requests.get", return_value=mock_response),
patch("litellm.proxy.client.cli.commands.auth.save_token") as mock_save,
patch("litellm.proxy.client.cli.interface.show_commands"),
patch("litellm._uuid.uuid.uuid4", return_value="session-uuid-solo"),
):
result = self.runner.invoke(login, obj=mock_context.obj)
@@ -637,7 +673,7 @@ class TestCLIKeyRegenerationFlow:
call_args = mock_browser.call_args[0][0]
assert "https://test.example.com/sso/key/generate" in call_args
assert "source=litellm-cli" in call_args
assert "key=sk-session-uuid-solo" in call_args
assert "key=cli-session-uuid-solo" in call_args
# Verify JWT was saved
mock_save.assert_called_once()
tests/test_litellm/proxy/management_endpoints/test_ui_sso.py
+248 -65
View File
@@ -4,7 +4,6 @@ import os
import sys
from unittest.mock import AsyncMock, MagicMock, patch
import httpx
import pytest
from fastapi import HTTPException, Request
@@ -25,7 +24,6 @@ from litellm.proxy.management_endpoints.ui_sso import (
SSOAuthenticationHandler,
_setup_team_mappings,
_sync_user_role_from_jwt_role_map,
determine_role_from_groups,
normalize_email,
process_sso_jwt_access_token,
)
@@ -1471,13 +1469,13 @@ class TestAuthCallbackRouting:
from litellm.constants import LITELLM_CLI_SESSION_TOKEN_PREFIX
# Test CLI state detection logic
cli_state = f"{LITELLM_CLI_SESSION_TOKEN_PREFIX}:sk-test123"
cli_state = f"{LITELLM_CLI_SESSION_TOKEN_PREFIX}:cli-test1234567890"
# This mimics the logic in auth_callback
if cli_state and cli_state.startswith(f"{LITELLM_CLI_SESSION_TOKEN_PREFIX}:"):
# Extract the key ID from the state
# Extract the login ID from the state
key_id = cli_state.split(":", 1)[1]
assert key_id == "sk-test123"
assert key_id == "cli-test1234567890"
else:
assert False, "CLI state should have been detected"
@@ -1510,13 +1508,13 @@ class TestGoogleLoginCLIIntegration:
# Test the CLI state generation logic used in google_login
source = "litellm-cli"
key = "sk-test123"
key = "cli-test1234567890"
cli_state = SSOAuthenticationHandler._get_cli_state(source=source, key=key)
assert cli_state is not None
assert cli_state.startswith("litellm-session-token:")
assert "sk-test123" in cli_state
assert "cli-test1234567890" in cli_state
def test_google_login_no_cli_state_when_missing_params(self):
"""Test that google_login doesn't generate CLI state when CLI parameters are missing"""
@@ -1526,8 +1524,8 @@ class TestGoogleLoginCLIIntegration:
test_cases = [
(None, None),
("litellm-cli", None),
(None, "sk-test123"),
("wrong-source", "sk-test123"),
(None, "cli-test1234567890"),
("wrong-source", "cli-test1234567890"),
]
for source, key in test_cases:
@@ -1634,19 +1632,19 @@ class TestSSOStateHandling:
from litellm.proxy.management_endpoints.ui_sso import SSOAuthenticationHandler
state = SSOAuthenticationHandler._get_cli_state(
source="litellm-cli", key="sk-test123"
source="litellm-cli", key="cli-test1234567890"
)
assert state is not None
assert state.startswith("litellm-session-token:")
assert "sk-test123" in state
assert "cli-test1234567890" in state
def test_get_cli_state_invalid_source(self):
"""Test generating CLI state with invalid source"""
from litellm.proxy.management_endpoints.ui_sso import SSOAuthenticationHandler
state = SSOAuthenticationHandler._get_cli_state(
source="invalid_source", key="sk-test123"
source="invalid_source", key="cli-test1234567890"
)
assert state is None
@@ -1663,40 +1661,40 @@ class TestSSOStateHandling:
"""Test generating CLI state without source"""
from litellm.proxy.management_endpoints.ui_sso import SSOAuthenticationHandler
state = SSOAuthenticationHandler._get_cli_state(source=None, key="sk-test123")
state = SSOAuthenticationHandler._get_cli_state(
source=None, key="cli-test1234567890"
)
assert state is None
def test_get_cli_state_with_existing_key(self):
"""Test generating CLI state with existing_key embedded in state parameter"""
def test_get_cli_state_ignores_existing_key(self):
"""Test CLI state does not embed an existing key"""
from litellm.proxy.management_endpoints.ui_sso import SSOAuthenticationHandler
state = SSOAuthenticationHandler._get_cli_state(
source="litellm-cli",
key="sk-new-key-123",
key="cli-new-key-1234567890",
existing_key="sk-existing-key-456",
)
assert state is not None
assert state.startswith("litellm-session-token:")
assert "sk-new-key-123" in state
assert "sk-existing-key-456" in state
# Verify the format: {PREFIX}:{key}:{existing_key}
assert state == "litellm-session-token:sk-new-key-123:sk-existing-key-456"
assert "cli-new-key-1234567890" in state
assert "sk-existing-key-456" not in state
assert state == "litellm-session-token:cli-new-key-1234567890"
def test_get_cli_state_without_existing_key(self):
"""Test generating CLI state without existing_key"""
from litellm.proxy.management_endpoints.ui_sso import SSOAuthenticationHandler
state = SSOAuthenticationHandler._get_cli_state(
source="litellm-cli", key="sk-new-key-789", existing_key=None
source="litellm-cli", key="cli-new-key-789123456", existing_key=None
)
assert state is not None
assert state.startswith("litellm-session-token:")
assert "sk-new-key-789" in state
# Verify the format: {PREFIX}:{key} (no third part)
assert state == "litellm-session-token:sk-new-key-789"
assert "cli-new-key-789123456" in state
assert state == "litellm-session-token:cli-new-key-789123456"
assert state.count(":") == 1 # Only one colon separator
@@ -1708,44 +1706,37 @@ class TestStateRouting:
from litellm.constants import LITELLM_CLI_SESSION_TOKEN_PREFIX
# Test CLI state format
cli_state = f"{LITELLM_CLI_SESSION_TOKEN_PREFIX}:sk-test123"
cli_state = f"{LITELLM_CLI_SESSION_TOKEN_PREFIX}:cli-test1234567890"
assert cli_state.startswith(f"{LITELLM_CLI_SESSION_TOKEN_PREFIX}:")
# Test extraction of key from state
key_id = cli_state.split(":", 1)[1]
assert key_id == "sk-test123"
assert key_id == "cli-test1234567890"
def test_cli_state_parsing_with_existing_key(self):
"""Test parsing CLI state with existing_key embedded"""
def test_cli_state_parsing_uses_single_login_id(self):
"""Test parsing CLI state with a single login ID"""
from litellm.constants import LITELLM_CLI_SESSION_TOKEN_PREFIX
# State format: {PREFIX}:{key}:{existing_key}
cli_state = (
f"{LITELLM_CLI_SESSION_TOKEN_PREFIX}:sk-new-key-456:sk-existing-key-789"
)
cli_state = f"{LITELLM_CLI_SESSION_TOKEN_PREFIX}:cli-new-key-456123"
# Parse as done in auth_callback
state_parts = cli_state.split(":", 2) # Split into max 3 parts
state_parts = cli_state.split(":", 1)
key_id = state_parts[1] if len(state_parts) > 1 else None
existing_key = state_parts[2] if len(state_parts) > 2 else None
assert key_id == "sk-new-key-456"
assert existing_key == "sk-existing-key-789"
assert key_id == "cli-new-key-456123"
def test_cli_state_parsing_without_existing_key(self):
"""Test parsing CLI state without existing_key"""
def test_cli_state_parsing_without_extra_segments(self):
"""Test parsing CLI state uses a single login ID"""
from litellm.constants import LITELLM_CLI_SESSION_TOKEN_PREFIX
# State format: {PREFIX}:{key}
cli_state = f"{LITELLM_CLI_SESSION_TOKEN_PREFIX}:sk-new-key-999"
cli_state = f"{LITELLM_CLI_SESSION_TOKEN_PREFIX}:cli-new-key-999123"
# Parse as done in auth_callback
state_parts = cli_state.split(":", 2) # Split into max 3 parts
state_parts = cli_state.split(":", 1)
key_id = state_parts[1] if len(state_parts) > 1 else None
existing_key = state_parts[2] if len(state_parts) > 2 else None
assert key_id == "sk-new-key-999"
assert existing_key is None
assert key_id == "cli-new-key-999123"
def test_non_cli_state_detection(self):
"""Test detection of non-CLI state parameters"""
@@ -2007,6 +1998,107 @@ class TestCustomUISSO:
class TestCLIKeyRegenerationFlow:
"""Test the end-to-end CLI key regeneration flow"""
@pytest.mark.asyncio
async def test_cli_sso_start_creates_bound_flow(self):
"""Test CLI SSO start creates a polling secret bound flow"""
from litellm.proxy.management_endpoints.ui_sso import (
_hash_cli_sso_secret,
_normalize_cli_sso_user_code,
cli_sso_start,
)
mock_cache = MagicMock()
with patch("litellm.proxy.proxy_server.user_api_key_cache", mock_cache):
result = await cli_sso_start()
assert result["login_id"].startswith("cli-")
assert result["poll_secret"]
assert result["user_code"]
mock_cache.set_cache.assert_called_once()
flow_data = mock_cache.set_cache.call_args.kwargs["value"]
assert flow_data["poll_secret_hash"] == _hash_cli_sso_secret(
result["poll_secret"]
)
assert flow_data["user_code_hash"] == _hash_cli_sso_secret(
_normalize_cli_sso_user_code(result["user_code"])
)
assert flow_data["poll_secret_hash"] != result["poll_secret"]
assert flow_data["user_code_hash"] != result["user_code"]
@pytest.mark.asyncio
async def test_cli_sso_complete_verifies_user_code(self):
"""Test CLI SSO complete marks a session as verified"""
from litellm.proxy.management_endpoints.ui_sso import (
_hash_cli_sso_secret,
_normalize_cli_sso_user_code,
cli_sso_complete,
)
mock_request = MagicMock(spec=Request)
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 = {
"poll_secret_hash": _hash_cli_sso_secret("poll-secret"),
"user_code_hash": _hash_cli_sso_secret(
_normalize_cli_sso_user_code("ABCD-EFGH")
),
"browser_complete_token_hash": _hash_cli_sso_secret("browser-token"),
"sso_complete": True,
"user_code_verified": False,
"session_data": {"user_id": "test-user-123"},
}
with (
patch("litellm.proxy.proxy_server.user_api_key_cache", mock_cache),
patch(
"litellm.proxy.common_utils.html_forms.cli_sso_success.render_cli_sso_success_page",
return_value="<html>Success</html>",
),
):
result = await cli_sso_complete(
request=mock_request, login_id="cli-session-4567890"
)
assert result.status_code == 200
flow_data = mock_cache.set_cache.call_args.kwargs["value"]
assert flow_data["user_code_verified"] is True
@pytest.mark.asyncio
async def test_cli_sso_complete_requires_callback_token(self):
"""Test CLI SSO complete requires the callback-delivered token"""
from litellm.proxy.management_endpoints.ui_sso import (
_hash_cli_sso_secret,
_normalize_cli_sso_user_code,
cli_sso_complete,
)
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 = {
"poll_secret_hash": _hash_cli_sso_secret("poll-secret"),
"user_code_hash": _hash_cli_sso_secret(
_normalize_cli_sso_user_code("ABCD-EFGH")
),
"browser_complete_token_hash": _hash_cli_sso_secret("browser-token"),
"sso_complete": True,
"user_code_verified": False,
"session_data": {"user_id": "test-user-123"},
}
with patch("litellm.proxy.proxy_server.user_api_key_cache", mock_cache):
with pytest.raises(HTTPException) as exc_info:
await cli_sso_complete(
request=mock_request, login_id="cli-session-4567890"
)
assert exc_info.value.status_code == 400
mock_cache.set_cache.assert_not_called()
@pytest.mark.asyncio
async def test_cli_sso_callback_stores_session(self):
"""Test CLI SSO callback stores session data in cache for JWT generation"""
@@ -2017,7 +2109,7 @@ class TestCLIKeyRegenerationFlow:
mock_request = MagicMock(spec=Request)
# Test data
session_key = "sk-session-456"
session_key = "cli-session-4567890"
# Mock user info
mock_user_info = LiteLLM_UserTable(
@@ -2032,6 +2124,16 @@ class TestCLIKeyRegenerationFlow:
# Mock cache
mock_cache = MagicMock()
mock_cache.get_cache.return_value = {
"poll_secret_hash": "poll-secret-hash",
"user_code_hash": "user-code-hash",
"sso_complete": False,
"user_code_verified": False,
"session_data": None,
}
mock_request.url_for.return_value = (
"https://test.example.com/sso/cli/complete/cli-session-4567890"
)
with (
patch(
@@ -2049,7 +2151,6 @@ class TestCLIKeyRegenerationFlow:
result = await cli_sso_callback(
request=mock_request,
key=session_key,
existing_key=None,
result=mock_sso_result,
)
@@ -2062,14 +2163,18 @@ class TestCLIKeyRegenerationFlow:
assert session_key in call_args.kwargs["key"]
# Verify session data structure
session_data = call_args.kwargs["value"]
flow_data = call_args.kwargs["value"]
session_data = flow_data["session_data"]
assert flow_data["sso_complete"] is True
assert flow_data["user_code_verified"] is False
assert isinstance(flow_data["browser_complete_token_hash"], str)
assert session_data["user_id"] == "test-user-123"
assert session_data["user_role"] == "internal_user"
assert session_data["teams"] == ["team1", "team2"]
assert session_data["models"] == ["gpt-4"]
# Verify TTL
assert call_args.kwargs["ttl"] == 600 # 10 minutes
assert call_args.kwargs["ttl"] == 600
assert result.status_code == 200
# Verify response contains success message (response is HTML)
@@ -2078,10 +2183,13 @@ class TestCLIKeyRegenerationFlow:
@pytest.mark.asyncio
async def test_cli_poll_key_returns_teams_for_selection(self):
"""Test CLI poll endpoint returns teams for user selection when multiple teams exist"""
from litellm.proxy.management_endpoints.ui_sso import cli_poll_key
from litellm.proxy.management_endpoints.ui_sso import (
_hash_cli_sso_secret,
cli_poll_key,
)
# Test data
session_key = "sk-session-789"
session_key = "cli-session-789123"
session_data = {
"user_id": "test-user-456",
"user_role": "internal_user",
@@ -2091,11 +2199,20 @@ class TestCLIKeyRegenerationFlow:
# Mock cache
mock_cache = MagicMock()
mock_cache.get_cache.return_value = session_data
mock_cache.get_cache.return_value = {
"poll_secret_hash": _hash_cli_sso_secret("poll-secret"),
"sso_complete": True,
"user_code_verified": True,
"session_data": session_data,
}
with patch("litellm.proxy.proxy_server.user_api_key_cache", mock_cache):
# Act - First poll without team_id
result = await cli_poll_key(key_id=session_key, team_id=None)
result = await cli_poll_key(
key_id=session_key,
team_id=None,
x_litellm_cli_poll_secret="poll-secret",
)
# Assert - should return teams list for selection
assert result["status"] == "ready"
@@ -2108,16 +2225,72 @@ class TestCLIKeyRegenerationFlow:
mock_cache.delete_cache.assert_not_called()
@pytest.mark.asyncio
async def test_auth_callback_routes_to_cli_with_existing_key(self):
"""Test that auth_callback properly routes CLI requests and extracts existing_key from state parameter"""
async def test_cli_poll_key_requires_poll_secret(self):
"""Test CLI poll endpoint rejects callers without the polling secret"""
from litellm.proxy.management_endpoints.ui_sso import (
_hash_cli_sso_secret,
cli_poll_key,
)
mock_cache = MagicMock()
mock_cache.get_cache.return_value = {
"poll_secret_hash": _hash_cli_sso_secret("poll-secret"),
"sso_complete": True,
"user_code_verified": True,
"session_data": {
"user_id": "test-user-456",
"user_role": "internal_user",
"teams": [],
"models": ["gpt-4"],
},
}
with patch("litellm.proxy.proxy_server.user_api_key_cache", mock_cache):
with pytest.raises(HTTPException) as exc_info:
await cli_poll_key(key_id="cli-session-789123", team_id=None)
assert exc_info.value.status_code == 403
@pytest.mark.asyncio
async def test_cli_poll_key_waits_for_user_code_verification(self):
"""Test CLI poll endpoint stays pending until user code verification"""
from litellm.proxy.management_endpoints.ui_sso import (
_hash_cli_sso_secret,
cli_poll_key,
)
mock_cache = MagicMock()
mock_cache.get_cache.return_value = {
"poll_secret_hash": _hash_cli_sso_secret("poll-secret"),
"sso_complete": True,
"user_code_verified": False,
"session_data": {
"user_id": "test-user-456",
"user_role": "internal_user",
"teams": [],
"models": ["gpt-4"],
},
}
with patch("litellm.proxy.proxy_server.user_api_key_cache", mock_cache):
result = await cli_poll_key(
key_id="cli-session-789123",
team_id=None,
x_litellm_cli_poll_secret="poll-secret",
)
assert result == {"status": "pending"}
@pytest.mark.asyncio
async def test_auth_callback_routes_to_cli(self):
"""Test that auth_callback properly routes CLI requests"""
from litellm.constants import LITELLM_CLI_SESSION_TOKEN_PREFIX
from litellm.proxy.management_endpoints.ui_sso import auth_callback
# Mock request (no query params needed - existing_key is in state)
# Mock request
mock_request = MagicMock(spec=Request)
# CLI state with existing_key embedded: {PREFIX}:{key}:{existing_key}
cli_state = f"{LITELLM_CLI_SESSION_TOKEN_PREFIX}:sk-new-session-key-456:sk-existing-cli-key-123"
cli_state = f"{LITELLM_CLI_SESSION_TOKEN_PREFIX}:cli-new-session-key-456"
# Mock the CLI callback and required proxy server components
mock_result = {"user_id": "test-user", "email": "test@example.com"}
@@ -2142,16 +2315,14 @@ class TestCLIKeyRegenerationFlow:
# Act
await auth_callback(request=mock_request, state=cli_state)
# Assert - existing_key should be extracted from state parameter
mock_cli_callback.assert_called_once_with(
request=mock_request,
key="sk-new-session-key-456",
existing_key="sk-existing-cli-key-123",
key="cli-new-session-key-456",
result=mock_result,
)
def test_get_redirect_url_does_not_include_existing_key_in_url(self):
"""Test that redirect URL generation does NOT include existing_key in URL (uses state parameter instead)"""
"""Test that redirect URL generation does NOT include existing_key in URL"""
from litellm.proxy.management_endpoints.ui_sso import SSOAuthenticationHandler
# Mock request
@@ -2194,10 +2365,13 @@ class TestCLIKeyRegenerationFlow:
async def test_cli_poll_key_generates_jwt_with_team(self):
"""Test CLI poll endpoint generates JWT when team_id is provided"""
from litellm.proxy._types import LiteLLM_UserTable
from litellm.proxy.management_endpoints.ui_sso import cli_poll_key
from litellm.proxy.management_endpoints.ui_sso import (
_hash_cli_sso_secret,
cli_poll_key,
)
# Test data
session_key = "sk-session-999"
session_key = "cli-session-999123"
selected_team = "team-b"
session_data = {
"user_id": "test-user-789",
@@ -2217,7 +2391,12 @@ class TestCLIKeyRegenerationFlow:
# Mock cache
mock_cache = MagicMock()
mock_cache.get_cache.return_value = session_data
mock_cache.get_cache.return_value = {
"poll_secret_hash": _hash_cli_sso_secret("poll-secret"),
"sso_complete": True,
"user_code_verified": True,
"session_data": session_data,
}
mock_jwt_token = "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.test.token"
@@ -2235,7 +2414,11 @@ class TestCLIKeyRegenerationFlow:
)
# Act - Second poll with team_id
result = await cli_poll_key(key_id=session_key, team_id=selected_team)
result = await cli_poll_key(
key_id=session_key,
team_id=selected_team,
x_litellm_cli_poll_secret="poll-secret",
)
# Assert - should return JWT
assert result["status"] == "ready"
@@ -2901,7 +3084,7 @@ class TestGetGenericSSORedirectParams:
from litellm.proxy.management_endpoints.ui_sso import SSOAuthenticationHandler
# Arrange
cli_state = "litellm-session-token:sk-test123"
cli_state = "litellm-session-token:cli-test1234567890"
with patch.dict(os.environ, {"GENERIC_CLIENT_STATE": "env_state_value"}):
# Act