diff --git a/litellm/constants.py b/litellm/constants.py
index a0e99dd16b..e2b7c864d4 100644
--- a/litellm/constants.py
+++ b/litellm/constants.py
@@ -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(
diff --git a/litellm/proxy/client/README.md b/litellm/proxy/client/README.md
index 5dcc88cacb..adf562d69c 100644
--- a/litellm/proxy/client/README.md
+++ b/litellm/proxy/client/README.md
@@ -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
diff --git a/litellm/proxy/client/cli/commands/auth.py b/litellm/proxy/client/cli/commands/auth.py
index aeb59e78a5..e9b370e4c0 100644
--- a/litellm/proxy/client/cli/commands/auth.py
+++ b/litellm/proxy/client/cli/commands/auth.py
@@ -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"]
diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py
index 46e7963da7..5485d618d5 100644
--- a/litellm/proxy/management_endpoints/ui_sso.py
+++ b/litellm/proxy/management_endpoints/ui_sso.py
@@ -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"""
+
+
+
+ LiteLLM CLI Login
+
+
+
+
+ Complete CLI Login
+ Enter the verification code shown in your terminal to finish this login.
+
+
+
+
+ """
+
+
+@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
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 45d55a8d06..f7cb4d72d9 100644
--- a/tests/test_litellm/proxy/client/cli/test_auth_commands.py
+++ b/tests/test_litellm/proxy/client/cli/test_auth_commands.py
@@ -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()
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 eecfcaa035..b4c843d0b8 100644
--- a/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py
+++ b/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py
@@ -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="Success",
+ ),
+ ):
+ 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