From bdb1e16dcf99026eba4cda0fa89dd34c736e1cae Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Mon, 17 Nov 2025 19:47:29 -0800 Subject: [PATCH] [Feat] AI Gateway Auth - Allow using JWTs for signing in with Proxy CLI (#16756) * fix auth * get_cli_jwt_auth_token * fix linting * test fixes * docs * test fixes * fix refactor --- docs/my-website/docs/proxy/cli_sso.md | 22 ++ docs/my-website/docs/proxy/management_cli.md | 19 ++ litellm/constants.py | 2 + litellm/proxy/auth/auth_checks.py | 51 +++ litellm/proxy/client/cli/commands/auth.py | 300 +++++++++++------- litellm/proxy/management_endpoints/ui_sso.py | 221 +++++++------ .../proxy/client/cli/test_auth_commands.py | 114 ++++--- .../test_callback_management_endpoints.py | 6 +- .../proxy/management_endpoints/test_ui_sso.py | 189 +++++++---- 9 files changed, 607 insertions(+), 317 deletions(-) diff --git a/docs/my-website/docs/proxy/cli_sso.md b/docs/my-website/docs/proxy/cli_sso.md index f7669d6a25..cde6bf266d 100644 --- a/docs/my-website/docs/proxy/cli_sso.md +++ b/docs/my-website/docs/proxy/cli_sso.md @@ -9,6 +9,26 @@ Use the litellm cli to authenticate to the LiteLLM Gateway. This is great if you ## Usage +### Prerequisites - Start LiteLLM Proxy with Beta Flag + +:::warning[Beta Feature - Required] + +CLI SSO Authentication is currently in beta. You must set this environment variable **when starting up your LiteLLM Proxy**: + +```bash +export EXPERIMENTAL_UI_LOGIN="True" +litellm --config config.yaml +``` + +Or add it to your proxy startup command: + +```bash +EXPERIMENTAL_UI_LOGIN="True" litellm --config config.yaml +``` + +::: + +### Steps 1. **Install the CLI** @@ -33,6 +53,8 @@ Use the litellm cli to authenticate to the LiteLLM Gateway. This is great if you 2. **Set up environment variables** + On your local machine, set the proxy URL: + ```bash export LITELLM_PROXY_URL=http://localhost:4000 ``` diff --git a/docs/my-website/docs/proxy/management_cli.md b/docs/my-website/docs/proxy/management_cli.md index 9ecc2ae8a3..23a5684210 100644 --- a/docs/my-website/docs/proxy/management_cli.md +++ b/docs/my-website/docs/proxy/management_cli.md @@ -67,7 +67,26 @@ For an indepth guide, see [CLI Authentication](./cli_sso). ::: +### Prerequisites +:::warning[Beta Feature - Required Environment Variable] + +CLI SSO Authentication is currently in beta. You must set this environment variable **when starting up your LiteLLM Proxy**: + +```bash +export EXPERIMENTAL_UI_LOGIN="True" +litellm --config config.yaml +``` + +Or add it to your proxy startup command: + +```bash +EXPERIMENTAL_UI_LOGIN="True" litellm --config config.yaml +``` + +::: + +### Steps 1. **Set up the proxy URL** diff --git a/litellm/constants.py b/litellm/constants.py index 5c4198ef1e..b90f36ae96 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1044,6 +1044,8 @@ LITELLM_PROXY_ADMIN_NAME = "default_user_id" ########################### CLI SSO AUTHENTICATION CONSTANTS ########################### LITELLM_CLI_SOURCE_IDENTIFIER = "litellm-cli" LITELLM_CLI_SESSION_TOKEN_PREFIX = "litellm-session-token" +CLI_SSO_SESSION_CACHE_KEY_PREFIX = "cli_sso_session" +CLI_JWT_TOKEN_NAME = "cli-jwt-token" ########################### DB CRON JOB NAMES ########################### DB_SPEND_UPDATE_JOB_NAME = "db_spend_update_job" diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index bfcf51a91d..04554aeb32 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -1352,6 +1352,57 @@ class ExperimentalUIJWTToken: return encrypt_value_helper(valid_token.model_dump_json(exclude_none=True)) + @staticmethod + def get_cli_jwt_auth_token( + user_info: LiteLLM_UserTable, team_id: Optional[str] = None + ) -> str: + """ + Generate a JWT token for CLI authentication with 24-hour expiration. + + Args: + user_info: User information from the database + team_id: Team ID for the user (optional, uses user's team if available) + + Returns: + Encrypted JWT token string + """ + from datetime import timedelta + + from litellm.constants import CLI_JWT_TOKEN_NAME + from litellm.proxy.common_utils.encrypt_decrypt_utils import ( + encrypt_value_helper, + ) + + if user_info.user_role is None: + raise Exception("User role is required for CLI JWT login") + + # Calculate expiration time (24 hours from now - matching old CLI key behavior) + expiration_time = get_utc_datetime() + timedelta(hours=24) + + # Format the expiration time as ISO 8601 string + expires = expiration_time.strftime("%Y-%m-%dT%H:%M:%S.%f")[:-3] + "+00:00" + + # Use provided team_id, or fall back to user's teams if available + _team_id = team_id + if _team_id is None and hasattr(user_info, "teams") and user_info.teams: + # Use first team if user has teams + _team_id = user_info.teams[0] if len(user_info.teams) > 0 else None + + valid_token = UserAPIKeyAuth( + token=CLI_JWT_TOKEN_NAME, + key_name=CLI_JWT_TOKEN_NAME, + key_alias=CLI_JWT_TOKEN_NAME, + max_budget=litellm.max_ui_session_budget, + expires=expires, + user_id=user_info.user_id, + team_id=_team_id, + models=user_info.models, + max_parallel_requests=None, + user_role=LitellmUserRoles(user_info.user_role), + ) + + return encrypt_value_helper(valid_token.model_dump_json(exclude_none=True)) + @staticmethod def get_key_object_from_ui_hash_key( hashed_token: str, diff --git a/litellm/proxy/client/cli/commands/auth.py b/litellm/proxy/client/cli/commands/auth.py index 1ac6ce3b13..49cbbed4b5 100644 --- a/litellm/proxy/client/cli/commands/auth.py +++ b/litellm/proxy/client/cli/commands/auth.py @@ -97,24 +97,6 @@ def display_teams_table(teams: List[Dict[str, Any]]) -> None: console.print(table) -def get_user_teams(base_url: str, api_key: str, user_id: str) -> List[Dict[str, Any]]: - """Fetch teams for the current user""" - from litellm.proxy.client import Client - - client = Client(base_url=base_url, api_key=api_key) - try: - response = client.teams.list_v2(user_id=user_id) - # Extract just the teams array from the paginated response - if isinstance(response, dict) and 'teams' in response: - return response['teams'] - else: - # Fallback in case the response structure is different - return response if isinstance(response, list) else [] - except Exception as e: - click.echo(f"āŒ Error fetching teams: {e}") - return [] - - def get_key_input(): """Get a single key input from the user (cross-platform)""" try: @@ -279,57 +261,172 @@ def prompt_team_selection_fallback(teams: List[Dict[str, Any]]) -> Optional[Dict return None -def update_key_with_team(base_url: str, api_key: str, team_id: str) -> bool: - """Update the API key to be associated with the selected team""" - from litellm.proxy._types import SpecialModelNames - from litellm.proxy.client import Client - - client = Client(base_url=base_url, api_key=api_key) - try: - client.keys.update(key=api_key, team_id=team_id, models=[SpecialModelNames.all_team_models.value]) - click.echo(f"āœ… Successfully assigned key to team: {team_id}") - return True - except requests.exceptions.HTTPError as e: - # Bubble up the response text for detailed error info - error_msg = e.response.text if e.response else str(e) - click.echo(f"āŒ Error updating key with team: {error_msg}") - return False - except Exception as e: - click.echo(f"āŒ Error updating key with team: {e}") - return False - - # Polling-based authentication - no local server needed -def _handle_team_assignment(base_url: str, api_key: str, user_id: str) -> None: - """Handle team fetching and assignment for the authenticated user.""" - click.echo("\n" + "="*60) - click.echo("šŸ“‹ Fetching your teams...") +def _poll_for_authentication( + base_url: str, key_id: str +) -> Optional[dict]: + """ + Poll the server for authentication completion and handle team selection. - teams = get_user_teams( - base_url=base_url, - api_key=api_key, - user_id=user_id, - ) + Returns: + Dictionary with authentication data if successful, None otherwise + """ + poll_url = f"{base_url}/sso/cli/poll/{key_id}" + timeout = 300 # 5 minute timeout + poll_interval = 2 # Poll every 2 seconds - if teams: - # Prompt for team selection (will display teams interactively) - selected_team = prompt_team_selection(teams) - - if selected_team: - team_id = selected_team.get('team_id') - if team_id: - click.echo(f"\nšŸ”„ Assigning your key to team: {selected_team.get('team_alias', team_id)}") - success = update_key_with_team(base_url, api_key, team_id) - if success: - click.echo(f"āœ… Your CLI key is now associated with team: {selected_team.get('team_alias', team_id)}") - click.echo(f"šŸŽÆ You can now access models: {', '.join(selected_team.get('models', ['All models']))}") - else: - click.echo("āš ļø Key assignment failed, but you can still use the CLI") + for attempt in range(timeout // poll_interval): + try: + response = requests.get(poll_url, timeout=10) + if response.status_code == 200: + data = response.json() + if data.get("status") == "ready": + # Check if we need team selection first + if data.get("requires_team_selection"): + # Server returned teams list without JWT - need to select team + teams = data.get("teams", []) + user_id = data.get("user_id") + + if teams and len(teams) > 1: + # User has multiple teams - let them select + jwt_with_team = _handle_team_selection_during_polling( + base_url=base_url, + key_id=key_id, + teams=teams + ) + + # Use the team-specific JWT if selection succeeded + if jwt_with_team: + return { + "api_key": jwt_with_team, + "user_id": user_id, + "teams": teams, + "team_id": None # Set by server in JWT + } + else: + # Selection failed or was skipped - poll again without team_id + click.echo("āš ļø Team selection skipped, retrying...") + continue + else: + # Shouldn't happen, but fallback + click.echo("āš ļø No teams available, retrying...") + continue + else: + # JWT is ready (single team or team already selected) + api_key = data.get("key") + user_id = data.get("user_id") + teams = data.get("teams", []) + team_id = data.get("team_id") + + # Show which team was assigned + if team_id and len(teams) == 1: + click.echo(f"\nāœ… Automatically assigned to team: {team_id}") + + if api_key: + return { + "api_key": api_key, + "user_id": user_id, + "teams": teams, + "team_id": team_id + } + elif data.get("status") == "pending": + # Still pending + if attempt % 10 == 0: # Show progress every 20 seconds + click.echo("Still waiting for authentication...") else: - click.echo("ā„¹ļø Continuing without team assignment. You can assign a team later using the CLI.") - else: + click.echo(f"Polling error: HTTP {response.status_code}") + + except requests.RequestException as e: + if attempt % 10 == 0: + click.echo(f"Connection error (will retry): {e}") + + time.sleep(poll_interval) + + # Timeout reached + return None + + +def _handle_team_selection_during_polling( + base_url: str, key_id: str, teams: List[str] +) -> Optional[str]: + """ + Handle team selection and re-poll with selected team_id. + + Args: + teams: List of team IDs (strings) + + Returns: + The JWT token with the selected team, or None if selection was skipped + """ + if not teams: click.echo("ā„¹ļø No teams found. You can create or join teams using the web interface.") + return None + + click.echo("\n" + "="*60) + click.echo("šŸ“‹ Select a team for your CLI session...") + + # Display teams as simple list since we only have IDs + console = Console() + table = Table(title="Available Teams") + table.add_column("Index", style="cyan", no_wrap=True) + table.add_column("Team ID", style="green") + + for i, team in enumerate(teams): + table.add_row(str(i + 1), team) + + console.print(table) + + # Simple selection + team_id: Optional[str] = None + while True: + try: + choice = click.prompt( + "\nSelect a team by entering the index number (or 'skip' to use first team)", + type=str + ).strip() + + if choice.lower() == 'skip': + team_id = teams[0] if teams else None + break + + index = int(choice) - 1 + if 0 <= index < len(teams): + team_id = teams[index] + break + else: + click.echo(f"āŒ Invalid selection. Please enter a number between 1 and {len(teams)}") + except ValueError: + click.echo("āŒ Invalid input. Please enter a number or 'skip'") + except KeyboardInterrupt: + click.echo("\nāŒ Team selection cancelled.") + return None + + if not team_id: + click.echo("ā„¹ļø No team selected.") + return None + + click.echo(f"\nšŸ”„ Generating JWT for team: {team_id}") + + # Re-poll with team_id to get JWT with correct team + try: + poll_url = f"{base_url}/sso/cli/poll/{key_id}?team_id={team_id}" + response = requests.get(poll_url, timeout=10) + + if response.status_code == 200: + data = response.json() + if data.get("status") == "ready": + jwt_token = data.get("key") + if jwt_token: + click.echo(f"āœ… Successfully generated JWT for team: {team_id}") + return jwt_token + + click.echo(f"āŒ Failed to get JWT with team. Status: {response.status_code}") + return None + + except Exception as e: + click.echo(f"āŒ Error getting JWT with team: {e}") + return None @click.command(name="login") @@ -337,7 +434,6 @@ def _handle_team_assignment(base_url: str, api_key: str, user_id: str) -> None: 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 @@ -365,59 +461,37 @@ def login(ctx: click.Context): # Open browser webbrowser.open(sso_url) - # Poll for key creation + # Poll for authentication completion click.echo("Waiting for authentication...") - poll_url = f"{base_url}/sso/cli/poll/{key_id}" - timeout = 300 # 5 minute timeout - poll_interval = 2 # Poll every 2 seconds + auth_result = _poll_for_authentication(base_url=base_url, key_id=key_id) - for attempt in range(timeout // poll_interval): - try: - response = requests.get(poll_url, timeout=10) - if response.status_code == 200: - data = response.json() - if data.get("status") == "ready": - # Key is ready - save it - api_key = data.get("key") - if api_key: - # Save token data (simplified for CLI - we just need the key) - save_token({ - 'key': api_key, - 'user_id': 'cli-user', - 'user_email': 'unknown', - 'user_role': 'cli', - 'auth_header_name': 'Authorization', - 'jwt_token': '', - 'timestamp': time.time() - }) - - click.echo("āœ… Login successful!") - click.echo(f"API Key: {api_key[:20]}...") - click.echo("You can now use the CLI without specifying --api-key") - - # Handle team assignment - _handle_team_assignment(base_url, api_key, data.get("user_id")) - - # Show available commands after successful login - click.echo("\n" + "="*60) - show_commands() - return - elif response.status_code == 200: - # Still pending - if attempt % 10 == 0: # Show progress every 20 seconds - click.echo("Still waiting for authentication...") - else: - click.echo(f"Polling error: HTTP {response.status_code}") - - except requests.RequestException as e: - if attempt % 10 == 0: - click.echo(f"Connection error (will retry): {e}") + if auth_result: + api_key = auth_result["api_key"] + user_id = auth_result["user_id"] - time.sleep(poll_interval) - - click.echo("āŒ Authentication timed out. Please try again.") - return + # Save token data (simplified for CLI - we just need the key) + save_token({ + 'key': api_key, + 'user_id': user_id or 'cli-user', + 'user_email': 'unknown', + 'user_role': 'cli', + 'auth_header_name': 'Authorization', + 'jwt_token': '', + 'timestamp': time.time() + }) + + click.echo("\nāœ… Login successful!") + click.echo(f"JWT Token: {api_key[:20]}...") + click.echo("You can now use the CLI without specifying --api-key") + + # Show available commands after successful login + click.echo("\n" + "="*60) + show_commands() + return + else: + click.echo("āŒ Authentication timed out. Please try again.") + return except KeyboardInterrupt: click.echo("\nāŒ Authentication cancelled by user.") diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py index e5341d35f6..ec103d773b 100644 --- a/litellm/proxy/management_endpoints/ui_sso.py +++ b/litellm/proxy/management_endpoints/ui_sso.py @@ -707,71 +707,22 @@ async def auth_callback(request: Request, state: Optional[str] = None): # noqa: ) -async def _regenerate_cli_key( - existing_key: str, new_key: str, user_id: Optional[str] = None -) -> None: - """Regenerate an existing CLI key with a new token""" - from litellm.proxy._types import RegenerateKeyRequest, UserAPIKeyAuth - from litellm.proxy.management_endpoints.key_management_endpoints import ( - regenerate_key_fn, - ) - - verbose_proxy_logger.info(f"Regenerating existing CLI key: {existing_key}") - - admin_user_dict = UserAPIKeyAuth.get_litellm_cli_user_api_key_auth() - - regenerate_request = RegenerateKeyRequest( - key=existing_key, - new_key=new_key, - duration="24hr", - user_id=user_id, - ) - - await regenerate_key_fn( - key=existing_key, data=regenerate_request, user_api_key_dict=admin_user_dict - ) - - verbose_proxy_logger.info(f"Regenerated CLI key: {new_key}") - - -async def _create_new_cli_key( - key: str, - user_id: Optional[str] = None, -) -> None: - """Create a new CLI key""" - from litellm.proxy.management_endpoints.key_management_endpoints import ( - generate_key_helper_fn, - ) - - verbose_proxy_logger.info("Creating new CLI key") - - await generate_key_helper_fn( - request_type="key", - duration="24hr", - key_max_budget=litellm.max_ui_session_budget, - aliases={}, - config={}, - spend=0, - user_id=user_id, - table_name="key", - token=key, - ) - - verbose_proxy_logger.info(f"Created new CLI key: {key}") - - 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 - regenerates existing CLI key or creates new one""" + """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}" ) - from litellm.proxy.proxy_server import prisma_client + from litellm.proxy.proxy_server import ( + prisma_client, + proxy_logging_obj, + user_api_key_cache, + ) if not key or not key.startswith("sk-"): raise HTTPException( @@ -784,24 +735,60 @@ async def cli_sso_callback( status_code=500, detail=CommonProxyErrors.db_not_connected_error.value ) + if result is None: + raise HTTPException( + status_code=500, + detail="SSO authentication failed - no result returned from provider", + ) + + # After None check, cast to non-None type for type checker + result_non_none: Union[OpenID, dict] = cast(Union[OpenID, dict], result) + parsed_openid_result = SSOAuthenticationHandler._get_user_email_and_id_from_result( - result=result + result=result_non_none ) verbose_proxy_logger.debug(f"parsed_openid_result: {parsed_openid_result}") try: - if existing_key: - await _regenerate_cli_key( - existing_key=existing_key, - new_key=key, - user_id=parsed_openid_result.get("user_id"), - ) - else: - await _create_new_cli_key( - key=key, - user_id=parsed_openid_result.get("user_id"), + # Get full user info from DB + user_info = await get_user_info_from_db( + result=result_non_none, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + user_email=parsed_openid_result.get("user_email"), + user_defined_values=None, + alternate_user_id=parsed_openid_result.get("user_id"), + ) + + if user_info is None: + raise HTTPException( + 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: + teams = user_info.teams if isinstance(user_info.teams, list) else [] + + session_data = { + "user_id": user_info.user_id, + "user_role": user_info.user_role, + "models": user_info.models if hasattr(user_info, "models") else [], + "user_email": parsed_openid_result.get("user_email"), + "teams": teams, + } + + cache_key = f"{CLI_SSO_SESSION_CACHE_KEY_PREFIX}:{key}" + user_api_key_cache.set_cache(key=cache_key, value=session_data, ttl=600) + + 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 @@ -813,47 +800,103 @@ async def cli_sso_callback( return HTMLResponse(content=html_content, status_code=200) except Exception as e: - verbose_proxy_logger.error(f"Error with CLI key: {e}") + verbose_proxy_logger.error(f"Error with CLI SSO callback: {e}") raise HTTPException( - status_code=500, detail=f"Failed to process CLI key: {str(e)}" + status_code=500, detail=f"Failed to process CLI SSO: {str(e)}" ) @router.get("/sso/cli/poll/{key_id}", tags=["experimental"], include_in_schema=False) -async def cli_poll_key(key_id: str): - """CLI polling endpoint - checks if key exists in DB""" - from litellm.proxy._types import LiteLLM_VerificationToken - from litellm.proxy.proxy_server import prisma_client +async def cli_poll_key(key_id: str, team_id: Optional[str] = None): + """ + CLI polling endpoint - retrieves session from cache and generates JWT. + + Flow: + 1. First poll (no team_id): Returns teams list without generating JWT + 2. Second poll (with team_id): Generates JWT with selected team and deletes session + + Args: + key_id: The session key 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") - if prisma_client is None: - raise HTTPException( - status_code=500, detail=CommonProxyErrors.db_not_connected_error.value - ) - try: - # Check if key exists in database - from litellm.proxy.utils import hash_token + # 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) - hashed_token = hash_token(key_id) + if session_data: + user_teams = session_data.get("teams", []) + user_id = session_data["user_id"] + + verbose_proxy_logger.info( + f"CLI poll: user={user_id}, team_id={team_id}, user_teams={user_teams}, num_teams={len(user_teams)}" + ) + + # If no team_id provided and user has teams, return teams list for selection + # Don't generate JWT yet - let CLI select a team first + if team_id is None and len(user_teams) > 1: + verbose_proxy_logger.info( + f"Returning teams list for user {user_id} to select from: {user_teams}" + ) + return { + "status": "ready", + "user_id": user_id, + "teams": user_teams, + "requires_team_selection": True, + } + + # Validate team_id if provided + if team_id is not None: + if team_id not in user_teams: + raise HTTPException( + status_code=403, + detail=f"User does not belong to team: {team_id}. Available teams: {user_teams}", + ) + else: + # If no team_id provided and user has 0 or 1 team, use first team (or None) + team_id = user_teams[0] if len(user_teams) > 0 else None - key_obj = await prisma_client.db.litellm_verificationtoken.find_unique( - where={"token": hashed_token} - ) - key_obj = cast(LiteLLM_VerificationToken, key_obj) + # Create user object for JWT generation + user_info = LiteLLM_UserTable( + user_id=user_id, + user_role=session_data["user_role"], + models=session_data.get("models", []), + max_budget=litellm.max_ui_session_budget, + ) - if key_obj: - verbose_proxy_logger.info(f"CLI key found: {key_id}") - return {"status": "ready", "key": key_id, "user_id": key_obj.user_id} + # Generate CLI JWT on-demand (24hr expiration) + # Pass selected team_id to ensure JWT has correct team + jwt_token = ExperimentalUIJWTToken.get_cli_jwt_auth_token( + user_info=user_info, team_id=team_id + ) + + # Delete cache entry (single-use) + user_api_key_cache.delete_cache(key=cache_key) + + verbose_proxy_logger.info( + f"CLI JWT generated for user: {user_id}, team: {team_id}" + ) + return { + "status": "ready", + "key": jwt_token, + "user_id": user_id, + "team_id": team_id, + "teams": user_teams, + } else: return {"status": "pending"} except Exception as e: - verbose_proxy_logger.error(f"Error polling for CLI key: {e}") + verbose_proxy_logger.error(f"Error polling for CLI JWT: {e}") raise HTTPException( - status_code=500, detail=f"Error checking key status: {str(e)}" + status_code=500, detail=f"Error checking session status: {str(e)}" ) 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 fa4ec0e98d..6e374cea47 100644 --- a/tests/test_litellm/proxy/client/cli/test_auth_commands.py +++ b/tests/test_litellm/proxy/client/cli/test_auth_commands.py @@ -188,16 +188,19 @@ class TestLoginCommand: self.runner = CliRunner() def test_login_success(self): - """Test successful login flow""" + """Test successful login flow with single team (JWT generated immediately)""" mock_context = Mock() mock_context.obj = {"base_url": "https://test.example.com"} - # Mock the requests for successful authentication + # Mock the requests for successful authentication with single team mock_response = Mock() mock_response.status_code = 200 mock_response.json.return_value = { "status": "ready", - "key": "sk-test-api-key-123" + "key": "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.test.jwt", + "user_id": "test-user-123", + "team_id": "team-1", + "teams": ["team-1"] } with patch('webbrowser.open') as mock_browser, \ @@ -210,7 +213,7 @@ class TestLoginCommand: assert result.exit_code == 0 assert "āœ… Login successful!" in result.output - assert "API Key: sk-test-api-key-123" in result.output + assert "Automatically assigned to team: team-1" in result.output # Verify browser was opened with correct URL mock_browser.assert_called_once() @@ -218,11 +221,11 @@ class TestLoginCommand: assert "https://test.example.com/sso/key/generate" in call_args assert "sk-test-uuid-123" in call_args - # Verify token was saved + # Verify JWT was saved mock_save.assert_called_once() saved_data = mock_save.call_args[0][0] - assert saved_data['key'] == 'sk-test-api-key-123' - assert saved_data['user_id'] == 'cli-user' + assert saved_data['key'] == "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.test.jwt" + assert saved_data['user_id'] == 'test-user-123' # Verify commands were shown mock_show_commands.assert_called_once() @@ -444,98 +447,107 @@ class TestCLIKeyRegenerationFlow: """Setup for each test""" self.runner = CliRunner() - def test_login_with_existing_key_regeneration_flow(self): - """Test complete login flow when user has existing key - should regenerate it""" + def test_login_with_team_selection_flow(self): + """Test complete login flow when user has multiple teams - should prompt for selection""" mock_context = Mock() mock_context.obj = {"base_url": "https://test.example.com"} - # Mock existing stored key - existing_key = "sk-existing-key-123" - - # Mock successful regeneration response - mock_response = Mock() - mock_response.status_code = 200 - mock_response.json.return_value = { + # Mock first response - requires team selection + mock_first_response = Mock() + mock_first_response.status_code = 200 + mock_first_response.json.return_value = { "status": "ready", - "key": "sk-regenerated-key-456" # New regenerated key + "requires_team_selection": True, + "user_id": "test-user-456", + "teams": ["team-alpha", "team-beta", "team-gamma"] } + # Mock second response after team selection - JWT with selected team + mock_second_response = Mock() + mock_second_response.status_code = 200 + mock_second_response.json.return_value = { + "status": "ready", + "key": "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.team-beta.jwt", + "user_id": "test-user-456", + "team_id": "team-beta", + "teams": ["team-alpha", "team-beta", "team-gamma"] + } + + # Simulate user selecting team #2 (team-beta) with patch('webbrowser.open') as mock_browser, \ - patch('requests.get', return_value=mock_response) as mock_get, \ - patch('litellm.proxy.client.cli.commands.auth.get_stored_api_key', return_value=existing_key) as mock_get_stored, \ + patch('requests.get', side_effect=[mock_first_response, mock_second_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='new-session-uuid-789'): + patch('litellm._uuid.uuid.uuid4', return_value='session-uuid-456'), \ + patch('click.prompt', return_value='2'): # User selects index 2 result = self.runner.invoke(login, obj=mock_context.obj) assert result.exit_code == 0 assert "āœ… Login successful!" in result.output - assert "API Key: sk-regenerated-key-4..." in result.output + assert "team-beta" in result.output - # Verify existing key was retrieved - mock_get_stored.assert_called_once() - - # Verify browser was opened with correct URL including existing key + # Verify browser was opened 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 "source=litellm-cli" in call_args - assert "key=sk-new-session-uuid-789" in call_args - assert f"existing_key={existing_key}" in call_args - # Verify polling was done with correct session key - mock_get.assert_called() - # Check that the polling URL was called (should be the first call) - first_call_args = mock_get.call_args_list[0] - poll_url = first_call_args[0][0] - assert "sk-new-session-uuid-789" in poll_url + # Verify two polling requests were made + assert mock_get.call_count == 2 - # Verify regenerated key was saved + # 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 "team_id=" not in first_poll_url + + # Second poll should include team_id=team-beta + second_poll_url = mock_get.call_args_list[1][0][0] + assert "team_id=team-beta" in second_poll_url + + # Verify JWT was saved mock_save.assert_called_once() saved_data = mock_save.call_args[0][0] - assert saved_data['key'] == 'sk-regenerated-key-456' - assert saved_data['user_id'] == 'cli-user' + assert saved_data['key'] == "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.team-beta.jwt" + assert saved_data['user_id'] == 'test-user-456' mock_show_commands.assert_called_once() - def test_login_without_existing_key_creation_flow(self): - """Test complete login flow when user has no existing key - should create new one""" + def test_login_without_teams_flow(self): + """Test complete login flow when user has no teams - JWT generated without team""" mock_context = Mock() mock_context.obj = {"base_url": "https://test.example.com"} - # Mock no existing key + # Mock response with no teams mock_response = Mock() mock_response.status_code = 200 mock_response.json.return_value = { "status": "ready", - "key": "sk-new-created-key-789" + "key": "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.no-team.jwt", + "user_id": "test-user-solo", + "team_id": None, + "teams": [] } with patch('webbrowser.open') as mock_browser, \ patch('requests.get', return_value=mock_response), \ - patch('litellm.proxy.client.cli.commands.auth.get_stored_api_key', return_value=None) as mock_get_stored, \ 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='new-session-uuid-999'): + patch('litellm._uuid.uuid.uuid4', return_value='session-uuid-solo'): result = self.runner.invoke(login, obj=mock_context.obj) assert result.exit_code == 0 assert "āœ… Login successful!" in result.output - # Verify existing key check was done - mock_get_stored.assert_called_once() - - # Verify browser was opened with correct URL WITHOUT existing key + # Verify browser was opened 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 "source=litellm-cli" in call_args - assert "key=sk-new-session-uuid-999" in call_args - assert "existing_key=" not in call_args # Should not include existing_key param + assert "key=sk-session-uuid-solo" in call_args - # Verify new key was saved + # Verify JWT was saved mock_save.assert_called_once() saved_data = mock_save.call_args[0][0] - assert saved_data['key'] == 'sk-new-created-key-789' + assert saved_data['key'] == "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.no-team.jwt" + assert saved_data['user_id'] == 'test-user-solo' diff --git a/tests/test_litellm/proxy/management_endpoints/test_callback_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_callback_management_endpoints.py index c30d6cf089..af00ab3677 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_callback_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_callback_management_endpoints.py @@ -249,6 +249,10 @@ class TestCallbackManagementEndpoints: # Verify dynamic_params structure assert isinstance(first_config["dynamic_params"], dict) - + # Check if at least one callback has detailed parameter configuration + has_detailed_params = any( + config.get("dynamic_params") and len(config.get("dynamic_params", {})) > 0 + for config in response_data + ) assert has_detailed_params, "Expected at least one callback to have detailed parameter configuration" 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 5266b1b5b3..26b34b994c 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py +++ b/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py @@ -1392,22 +1392,41 @@ class TestCLIKeyRegenerationFlow: """Test the end-to-end CLI key regeneration flow""" @pytest.mark.asyncio - async def test_cli_sso_callback_regenerate_existing_key(self): - """Test CLI SSO callback regenerating an existing key""" + async def test_cli_sso_callback_stores_session(self): + """Test CLI SSO callback stores session data in cache for JWT generation""" + from litellm.proxy._types import LiteLLM_UserTable from litellm.proxy.management_endpoints.ui_sso import cli_sso_callback # Mock request mock_request = MagicMock(spec=Request) # Test data - existing_key = "sk-existing-key-123" - new_key = "sk-new-key-456" + session_key = "sk-session-456" + + # Mock user info + mock_user_info = LiteLLM_UserTable( + user_id="test-user-123", + user_role="internal_user", + teams=["team1", "team2"], + models=["gpt-4"] + ) - # Mock the regenerate helper function + # Mock SSO result + mock_sso_result = { + "user_email": "test@example.com", + "user_id": "test-user-123" + } + + # Mock cache + mock_cache = MagicMock() + with patch( - "litellm.proxy.management_endpoints.ui_sso._regenerate_cli_key" - ) as mock_regenerate, patch( + "litellm.proxy.management_endpoints.ui_sso.get_user_info_from_db", + return_value=mock_user_info + ), patch( "litellm.proxy.proxy_server.prisma_client", MagicMock() + ), 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", @@ -1415,46 +1434,65 @@ class TestCLIKeyRegenerationFlow: # Act result = await cli_sso_callback( - request=mock_request, key=new_key, existing_key=existing_key + request=mock_request, key=session_key, existing_key=None, result=mock_sso_result ) - # Assert - mock_regenerate.assert_called_once_with( - existing_key=existing_key, new_key=new_key, user_id=None - ) + # Assert - verify session was stored in cache + mock_cache.set_cache.assert_called_once() + call_args = mock_cache.set_cache.call_args + + # Verify cache key format + assert "cli_sso_session:" in call_args.kwargs["key"] + assert session_key in call_args.kwargs["key"] + + # Verify session data structure + session_data = call_args.kwargs["value"] + 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 result.status_code == 200 - assert "Success" in result.body.decode() + # Verify response contains success message (response is HTML) + assert result.body is not None @pytest.mark.asyncio - async def test_cli_sso_callback_create_new_key(self): - """Test CLI SSO callback creating a new key when no existing key provided""" - from litellm.proxy.management_endpoints.ui_sso import cli_sso_callback - - # Mock request - mock_request = MagicMock(spec=Request) + 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 # Test data - new_key = "sk-new-key-789" + session_key = "sk-session-789" + session_data = { + "user_id": "test-user-456", + "user_role": "internal_user", + "teams": ["team-a", "team-b", "team-c"], + "models": ["gpt-4"] + } - # Mock the create helper function + # Mock cache + mock_cache = MagicMock() + mock_cache.get_cache.return_value = session_data + with patch( - "litellm.proxy.management_endpoints.ui_sso._create_new_cli_key" - ) as mock_create, patch( - "litellm.proxy.proxy_server.prisma_client", MagicMock() - ), patch( - "litellm.proxy.common_utils.html_forms.cli_sso_success.render_cli_sso_success_page", - return_value="Success", + "litellm.proxy.proxy_server.user_api_key_cache", mock_cache ): - # Act - result = await cli_sso_callback( - request=mock_request, key=new_key, existing_key=None - ) + # Act - First poll without team_id + result = await cli_poll_key(key_id=session_key, team_id=None) - # Assert - mock_create.assert_called_once_with(key=new_key, user_id=None) - assert result.status_code == 200 - assert "Success" in result.body.decode() + # Assert - should return teams list for selection + assert result["status"] == "ready" + assert result["requires_team_selection"] is True + assert result["user_id"] == "test-user-456" + assert result["teams"] == ["team-a", "team-b", "team-c"] + assert "key" not in result # JWT should not be generated yet + + # Verify session was NOT deleted + mock_cache.delete_cache.assert_not_called() @pytest.mark.asyncio async def test_auth_callback_routes_to_cli_with_existing_key(self): @@ -1543,40 +1581,65 @@ class TestCLIKeyRegenerationFlow: assert "https://test.litellm.ai/sso/callback" == redirect_url @pytest.mark.asyncio - async def test_cli_sso_callback_regenerate_vs_create_flow(self): - """Test CLI SSO callback calls regenerate_key_fn when existing_key provided, generate_key_helper_fn when not""" - from litellm.proxy.management_endpoints.ui_sso import cli_sso_callback + 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 - mock_request = MagicMock(spec=Request) + # Test data + session_key = "sk-session-999" + selected_team = "team-b" + session_data = { + "user_id": "test-user-789", + "user_role": "internal_user", + "teams": ["team-a", "team-b", "team-c"], + "models": ["gpt-4"], + "user_email": "test@example.com" + } + + # Mock user info + mock_user_info = LiteLLM_UserTable( + user_id="test-user-789", + user_role="internal_user", + teams=["team-a", "team-b", "team-c"], + models=["gpt-4"] + ) + # Mock cache + mock_cache = MagicMock() + mock_cache.get_cache.return_value = session_data + + mock_jwt_token = "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.test.token" + with patch( - "litellm.proxy.management_endpoints.key_management_endpoints.regenerate_key_fn" - ) as mock_regenerate, patch( - "litellm.proxy.management_endpoints.key_management_endpoints.generate_key_helper_fn" - ) as mock_generate, patch( - "litellm.proxy._types.UserAPIKeyAuth.get_litellm_cli_user_api_key_auth" + "litellm.proxy.proxy_server.user_api_key_cache", mock_cache ), patch( - "litellm.proxy.proxy_server.prisma_client", MagicMock() - ), patch( - "litellm.proxy.common_utils.html_forms.cli_sso_success.render_cli_sso_success_page", - return_value="Success", - ): + "litellm.proxy.proxy_server.prisma_client" + ) as mock_prisma, patch( + "litellm.proxy.auth.auth_checks.ExperimentalUIJWTToken.get_cli_jwt_auth_token", + return_value=mock_jwt_token + ) as mock_get_jwt: + + # Mock the user lookup + mock_prisma.db.litellm_usertable.find_unique = AsyncMock(return_value=mock_user_info) - # Test regeneration path - await cli_sso_callback( - mock_request, key="sk-new-123", existing_key="sk-existing-456" - ) - mock_regenerate.assert_called_once() - mock_generate.assert_not_called() + # Act - Second poll with team_id + result = await cli_poll_key(key_id=session_key, team_id=selected_team) - # Reset mocks - mock_regenerate.reset_mock() - mock_generate.reset_mock() - - # Test creation path - await cli_sso_callback(mock_request, key="sk-new-789", existing_key=None) - mock_regenerate.assert_not_called() - mock_generate.assert_called_once() + # Assert - should return JWT + assert result["status"] == "ready" + assert result["key"] == mock_jwt_token + assert result["user_id"] == "test-user-789" + assert result["team_id"] == selected_team + assert result["teams"] == ["team-a", "team-b", "team-c"] + + # Verify JWT was generated with correct team + mock_get_jwt.assert_called_once() + jwt_call_args = mock_get_jwt.call_args + assert jwt_call_args.kwargs["team_id"] == selected_team + + # Verify session was deleted after JWT generation + mock_cache.delete_cache.assert_called_once() class TestGetAppRolesFromIdToken: