[Feat] Proxy CLI Auth - Allow re-using cli auth token (#14780)

* fix: cli auth with SSO okta

* fix: add LITTELM_CLI_SERVICE_ACCOUNT_NAME

* fix: get_litellm_cli_user_api_key_auth

* use existing_key CLI

* fix: use existing key

* test auth commands

* test_cli_sso_callback_regenerate_vs_create_flow

Author: ishaan-jaff

litellm/constants.py

@@ -947,6 +947,7 @@
     os.getenv("HEALTH_CHECK_TIMEOUT_SECONDS", 60)
 )  # 60 seconds
 LITTELM_INTERNAL_HEALTH_SERVICE_ACCOUNT_NAME = "litellm-internal-health-check"
+LITTELM_CLI_SERVICE_ACCOUNT_NAME = "litellm-cli"
 
 UI_SESSION_TOKEN_TEAM_ID = "litellm-dashboard"
 LITELLM_PROXY_ADMIN_NAME = "default_user_id"

litellm/proxy/_types.py

@@ -1915,6 +1915,22 @@ def get_litellm_internal_health_check_user_api_key_auth(cls) -> "UserAPIKeyAuth"
             key_alias=LITTELM_INTERNAL_HEALTH_SERVICE_ACCOUNT_NAME,
             team_alias=LITTELM_INTERNAL_HEALTH_SERVICE_ACCOUNT_NAME,
         )
+    
+    @classmethod
+    def get_litellm_cli_user_api_key_auth(cls) -> "UserAPIKeyAuth":
+        """
+        Returns a `UserAPIKeyAuth` object for the litellm internal health check service account.
+
+        This is used to track number of requests/spend for health check calls.
+        """
+        from litellm.constants import LITTELM_CLI_SERVICE_ACCOUNT_NAME
+
+        return cls(
+            api_key=LITTELM_CLI_SERVICE_ACCOUNT_NAME,
+            team_id=LITTELM_CLI_SERVICE_ACCOUNT_NAME,
+            key_alias=LITTELM_CLI_SERVICE_ACCOUNT_NAME,
+            team_alias=LITTELM_CLI_SERVICE_ACCOUNT_NAME,
+        )
 
 
 class UserInfoResponse(LiteLLMPydanticObjectBase):

litellm/proxy/client/cli/commands/auth.py

@@ -64,13 +64,20 @@ def login(ctx: click.Context):
 
     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}"
 
+        # If we have an existing key, include it so the server can regenerate it
+        if existing_key:
+            sso_url += f"&existing_key={existing_key}"
+        
         click.echo(f"Opening browser to: {sso_url}")
         click.echo("Please complete the SSO authentication in your browser...")
         click.echo(f"Session ID: {key_id}")

litellm/proxy/management_endpoints/ui_sso.py

@@ -114,7 +114,7 @@ def process_sso_jwt_access_token(
 
 @router.get("/sso/key/generate", tags=["experimental"], include_in_schema=False)
 async def google_login(
-    request: Request, source: Optional[str] = None, key: Optional[str] = None
+    request: Request, source: Optional[str] = None, key: Optional[str] = None, existing_key: Optional[str] = None
 ):  # noqa: PLR0915
     """
     Create Proxy API Keys using Google Workspace SSO. Requires setting PROXY_BASE_URL in .env
@@ -173,12 +173,14 @@ 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
     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
@@ -590,8 +592,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 from the state
         key_id = state.split(":", 1)[1]
-        verbose_proxy_logger.info(f"CLI SSO callback detected for key: {key_id}")
-        return await cli_sso_callback(request, key=key_id)
+        
+        # Get existing_key from query parameters if provided
+        existing_key = request.query_params.get("existing_key")
+            
+        verbose_proxy_logger.info(f"CLI SSO callback detected for key: {key_id}, existing_key: {existing_key}")
+        return await cli_sso_callback(request, key=key_id, existing_key=existing_key)
 
     from litellm.proxy._types import LiteLLM_JWTAuth
     from litellm.proxy.auth.handle_jwt import JWTHandler
@@ -678,13 +684,59 @@ async def auth_callback(request: Request, state: Optional[str] = None):  # noqa:
     )
 
 
-async def cli_sso_callback(request: Request, key: Optional[str] = None):
-    """CLI SSO callback - generates the key with pre-specified ID"""
-    verbose_proxy_logger.info(f"CLI SSO callback for key: {key}")
+async def _regenerate_cli_key(existing_key: str, new_key: str) -> 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"
+    )
+    
+    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) -> 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,
+        team_id="litellm-cli",
+        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):
+    """CLI SSO callback - regenerates existing CLI key or creates new one"""
+    verbose_proxy_logger.info(f"CLI SSO callback for key: {key}, existing_key: {existing_key}")
+
     from litellm.proxy.proxy_server import prisma_client
 
     if not key or not key.startswith("sk-"):
@@ -698,21 +750,11 @@ async def cli_sso_callback(request: Request, key: Optional[str] = None):
             status_code=500, detail=CommonProxyErrors.db_not_connected_error.value
         )
 
-    # Generate a simple key for CLI usage with the pre-specified key ID
     try:
-        await generate_key_helper_fn(
-            request_type="key",
-            duration="24hr",
-            key_max_budget=litellm.max_ui_session_budget,
-            aliases={},
-            config={},
-            spend=0,
-            team_id="litellm-cli",
-            table_name="key",
-            token=key,  # Use the pre-specified key ID
-        )
-
-        verbose_proxy_logger.info(f"Generated CLI key: {key}")
+        if existing_key:
+            await _regenerate_cli_key(existing_key, key)
+        else:
+            await _create_new_cli_key(key)
 
         # Return success page
         from fastapi.responses import HTMLResponse
@@ -725,8 +767,8 @@ async def cli_sso_callback(request: Request, key: Optional[str] = None):
         return HTMLResponse(content=html_content, status_code=200)
 
     except Exception as e:
-        verbose_proxy_logger.error(f"Error generating CLI key: {e}")
-        raise HTTPException(status_code=500, detail=f"Failed to generate key: {str(e)}")
+        verbose_proxy_logger.error(f"Error with CLI key: {e}")
+        raise HTTPException(status_code=500, detail=f"Failed to process CLI key: {str(e)}")
 
 
 @router.get("/sso/cli/poll/{key_id}", tags=["experimental"], include_in_schema=False)
@@ -993,20 +1035,51 @@ async def get_sso_login_redirect(
                 # or a cryptographicly signed state that we can verify stateless
                 # For simplification we are using a static state, this is not perfect but some
                 # SSO providers do not allow stateless verification
-                redirect_params = {}
-                state = os.getenv("GENERIC_CLIENT_STATE", None)
-
-                if state:
-                    redirect_params["state"] = state
-                elif "okta" in generic_authorization_endpoint:
-                    redirect_params["state"] = (
-                        uuid.uuid4().hex
-                    )  # set state param for okta - required
+                redirect_params = SSOAuthenticationHandler._get_generic_sso_redirect_params(
+                    state=state,
+                    generic_authorization_endpoint=generic_authorization_endpoint
+                )
+                
                 return await generic_sso.get_login_redirect(**redirect_params)  # type: ignore
         raise ValueError(
             "Unknown SSO provider. Please setup SSO with client IDs https://docs.litellm.ai/docs/proxy/admin_ui_sso"
         )
 
+    @staticmethod
+    def _get_generic_sso_redirect_params(
+        state: Optional[str] = None, 
+        generic_authorization_endpoint: Optional[str] = None
+    ) -> dict:
+        """
+        Get redirect parameters for Generic SSO with proper state priority handling.
+        
+        Priority order:
+        1. CLI state (if provided)
+        2. GENERIC_CLIENT_STATE environment variable
+        3. Generated UUID for Okta (if Okta endpoint detected)
+        
+        Args:
+            state: Optional state parameter (e.g., CLI state)
+            generic_authorization_endpoint: Authorization endpoint URL
+            
+        Returns:
+            dict: Redirect parameters for SSO login
+        """
+        redirect_params = {}
+        
+        if state:
+            # CLI state takes priority
+            # the litellm proxy cli sends the "state" parameter to the proxy server for auth. We should maintain the state parameter for the cli if it is provided
+            redirect_params["state"] = state
+        else:
+            generic_client_state = os.getenv("GENERIC_CLIENT_STATE", None)
+            if generic_client_state:
+                redirect_params["state"] = generic_client_state
+            elif generic_authorization_endpoint and "okta" in generic_authorization_endpoint:
+                redirect_params["state"] = uuid.uuid4().hex  # set state param for okta - required
+
+        return redirect_params
+
     @staticmethod
     def should_use_sso_handler(
         google_client_id: Optional[str] = None,
@@ -1025,6 +1098,7 @@ def should_use_sso_handler(
     def get_redirect_url_for_sso(
         request: Request,
         sso_callback_route: str,
+        existing_key: Optional[str] = None,
     ) -> str:
         """
         Get the redirect URL for SSO
@@ -1036,6 +1110,11 @@ def get_redirect_url_for_sso(
             redirect_url += sso_callback_route
         else:
             redirect_url += "/" + sso_callback_route
+            
+        # Append existing_key as query parameter if provided
+        if existing_key:
+            redirect_url += f"?existing_key={existing_key}"
+            
         return redirect_url
 
     @staticmethod
@@ -1218,7 +1297,7 @@ def _cast_and_deepcopy_litellm_default_team_params(
         return team_request
 
     @staticmethod
-    def _get_cli_state(source: Optional[str], key: Optional[str]) -> Optional[str]:
+    def _get_cli_state(source: Optional[str], key: Optional[str], existing_key: Optional[str] = None) -> Optional[str]:
         """
         Checks the request 'source' if a cli state token was passed in
 
@@ -1229,11 +1308,11 @@ def _get_cli_state(source: Optional[str], key: Optional[str]) -> Optional[str]:
             LITELLM_CLI_SOURCE_IDENTIFIER,
         )
 
-        return (
-            f"{LITELLM_CLI_SESSION_TOKEN_PREFIX}:{key}"
-            if source == LITELLM_CLI_SOURCE_IDENTIFIER and key
-            else None
-        )
+        if source == LITELLM_CLI_SOURCE_IDENTIFIER and key:
+            # Just use the key - existing_key will be passed separately via query params
+            return f"{LITELLM_CLI_SESSION_TOKEN_PREFIX}:{key}"
+        else:
+            return None
 
     @staticmethod
     async def get_redirect_response_from_openid(  # noqa: PLR0915