Merge branch 'BerriAI:main' into LangfuseUsageDetails

litellm/constants.py

@@ -947,6 +947,7 @@ HEALTH_CHECK_TIMEOUT_SECONDS = int(
     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 @@ class 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,6 +64,9 @@ 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())}"
     
@@ -71,6 +74,10 @@ def login(ctx: click.Context):
         # 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 @@ class SSOAuthenticationHandler:
                 # 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 @@ class SSOAuthenticationHandler:
     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 @@ class SSOAuthenticationHandler:
             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 @@ class SSOAuthenticationHandler:
         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 @@ class SSOAuthenticationHandler:
             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

tests/test_litellm/proxy/client/cli/test_auth_commands.py

@@ -435,3 +435,105 @@ class TestWhoamiCommand:
             assert "✅ Authenticated" in result.output
             # Should calculate age based on timestamp=0
             assert "Token age:" in result.output
+
+
+class TestCLIKeyRegenerationFlow:
+    """Test the end-to-end CLI key regeneration flow from CLI perspective"""
+
+    def setup_method(self):
+        """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"""
+        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 = {
+            "status": "ready",
+            "key": "sk-regenerated-key-456"  # New regenerated key
+        }
+        
+        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('litellm.proxy.client.cli.commands.auth.save_token') as mock_save, \
+             patch('litellm.proxy.client.cli.interface.show_commands') as mock_show_commands, \
+             patch('uuid.uuid4', return_value='new-session-uuid-789'):
+            
+            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-456" in result.output
+            
+            # Verify existing key was retrieved
+            mock_get_stored.assert_called_once()
+            
+            # Verify browser was opened with correct URL including existing key
+            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()
+            poll_url = mock_get.call_args[0][0]
+            assert "sk-new-session-uuid-789" in poll_url
+            
+            # Verify regenerated key 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'
+            
+            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"""
+        mock_context = Mock()
+        mock_context.obj = {"base_url": "https://test.example.com"}
+        
+        # Mock no existing key
+        mock_response = Mock()
+        mock_response.status_code = 200
+        mock_response.json.return_value = {
+            "status": "ready",
+            "key": "sk-new-created-key-789"
+        }
+        
+        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('uuid.uuid4', return_value='new-session-uuid-999'):
+            
+            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
+            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
+            
+            # Verify new key was saved
+            mock_save.assert_called_once()
+            saved_data = mock_save.call_args[0][0]
+            assert saved_data['key'] == 'sk-new-created-key-789'

tests/test_litellm/proxy/management_endpoints/test_ui_sso.py

@@ -1247,6 +1247,156 @@ class TestCustomUISSO:
                     assert result.status_code == 303
 
 
+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"""
+        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"
+        
+        # Mock the regenerate helper function
+        with patch("litellm.proxy.management_endpoints.ui_sso._regenerate_cli_key") as mock_regenerate, \
+             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="<html>Success</html>"):
+            
+            # Act
+            result = await cli_sso_callback(
+                request=mock_request,
+                key=new_key,
+                existing_key=existing_key
+            )
+            
+            # Assert
+            mock_regenerate.assert_called_once_with(existing_key, new_key)
+            assert result.status_code == 200
+            assert "Success" in result.body.decode()
+
+    @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)
+        
+        # Test data
+        new_key = "sk-new-key-789"
+        
+        # Mock the create helper function
+        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="<html>Success</html>"):
+            
+            # Act
+            result = await cli_sso_callback(
+                request=mock_request,
+                key=new_key,
+                existing_key=None
+            )
+            
+            # Assert
+            mock_create.assert_called_once_with(new_key)
+            assert result.status_code == 200
+            assert "Success" in result.body.decode()
+
+    @pytest.mark.asyncio
+    async def test_auth_callback_routes_to_cli_with_existing_key(self):
+        """Test that auth_callback properly routes CLI requests and preserves existing_key parameter"""
+        from litellm.constants import LITELLM_CLI_SESSION_TOKEN_PREFIX
+        from litellm.proxy.management_endpoints.ui_sso import auth_callback
+
+        # Mock request with existing_key query parameter
+        mock_request = MagicMock(spec=Request)
+        mock_request.query_params.get.return_value = "sk-existing-cli-key-123"
+        
+        # CLI state
+        cli_state = f"{LITELLM_CLI_SESSION_TOKEN_PREFIX}:sk-new-session-key-456"
+        
+        # Mock the CLI callback
+        with patch("litellm.proxy.management_endpoints.ui_sso.cli_sso_callback") as mock_cli_callback:
+            mock_cli_callback.return_value = MagicMock()
+            
+            # Act
+            await auth_callback(request=mock_request, state=cli_state)
+            
+            # Assert
+            mock_cli_callback.assert_called_once_with(
+                mock_request,
+                key="sk-new-session-key-456",
+                existing_key="sk-existing-cli-key-123"
+            )
+
+    def test_get_redirect_url_preserves_existing_key(self):
+        """Test that redirect URL generation preserves existing_key parameter"""
+        from litellm.proxy.management_endpoints.ui_sso import SSOAuthenticationHandler
+
+        # Mock request
+        mock_request = MagicMock()
+        mock_request.base_url = "https://test.litellm.ai/"
+        
+        with patch("litellm.proxy.utils.get_custom_url", return_value="https://test.litellm.ai"):
+            # Test with existing_key
+            redirect_url = SSOAuthenticationHandler.get_redirect_url_for_sso(
+                request=mock_request,
+                sso_callback_route="sso/callback",
+                existing_key="sk-existing-123"
+            )
+            
+            assert "https://test.litellm.ai/sso/callback?existing_key=sk-existing-123" == redirect_url
+
+    def test_get_redirect_url_without_existing_key(self):
+        """Test that redirect URL generation works without existing_key parameter"""
+        from litellm.proxy.management_endpoints.ui_sso import SSOAuthenticationHandler
+
+        # Mock request
+        mock_request = MagicMock()
+        mock_request.base_url = "https://test.litellm.ai/"
+        
+        with patch("litellm.proxy.utils.get_custom_url", return_value="https://test.litellm.ai"):
+            # Test without existing_key
+            redirect_url = SSOAuthenticationHandler.get_redirect_url_for_sso(
+                request=mock_request,
+                sso_callback_route="sso/callback"
+            )
+            
+            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
+        
+        mock_request = MagicMock(spec=Request)
+        
+        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"), \
+             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="<html>Success</html>"):
+            
+            # 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()
+            
+            # 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()
+
+
 class TestProcessSSOJWTAccessToken:
     """Test the process_sso_jwt_access_token helper function"""