feat: Add group delete and rm --group commands with workload management (#1174)

Author: amirejaz

cmd/thv/app/group.go

@@ -1,6 +1,8 @@
 package app
 
 import (
+	"bufio"
+	"context"
 	"fmt"
 	"os"
 	"sort"
@@ -10,6 +12,7 @@ import (
 	"github.com/spf13/cobra"
 
 	"github.com/stacklok/toolhive/pkg/groups"
+	"github.com/stacklok/toolhive/pkg/workloads"
 )
 
 var groupCmd = &cobra.Command{
@@ -33,6 +36,17 @@ var groupListCmd = &cobra.Command{
 	RunE:  groupListCmdFunc,
 }
 
+var groupRmCmd = &cobra.Command{
+	Use:   "rm [group-name]",
+	Short: "Remove a group and remove workloads from it",
+	Long: "Remove a group and remove all MCP servers from it. By default, this only removes the group " +
+		"membership from workloads without deleting them. Use --with-workloads to also delete the workloads. ",
+	Args: cobra.ExactArgs(1),
+	RunE: groupRmCmdFunc,
+}
+
+var withWorkloadsFlag bool
+
 func groupCreateCmdFunc(cmd *cobra.Command, args []string) error {
 	groupName := args[0]
 	ctx := cmd.Context()
@@ -90,7 +104,153 @@ func groupListCmdFunc(cmd *cobra.Command, _ []string) error {
 	return nil
 }
 
+func groupRmCmdFunc(cmd *cobra.Command, args []string) error {
+	groupName := args[0]
+	ctx := cmd.Context()
+
+	if strings.EqualFold(groupName, groups.DefaultGroup) {
+		return fmt.Errorf("cannot delete the %s group", groups.DefaultGroup)
+	}
+
+	groupManager, err := groups.NewManager()
+	if err != nil {
+		return fmt.Errorf("failed to create group manager: %w", err)
+	}
+
+	// Check if group exists
+	exists, err := groupManager.Exists(ctx, groupName)
+	if err != nil {
+		return fmt.Errorf("failed to check if group exists: %w", err)
+	}
+	if !exists {
+		return fmt.Errorf("group '%s' does not exist", groupName)
+	}
+
+	// Get all workloads in the group
+	groupWorkloads, err := groupManager.ListWorkloadsInGroup(ctx, groupName)
+	if err != nil {
+		return fmt.Errorf("failed to list workloads in group: %w", err)
+	}
+
+	// Show warning and get user confirmation
+	confirmed, err := showWarningAndGetConfirmation(groupName, groupWorkloads)
+	if err != nil {
+		return err
+	}
+
+	if !confirmed {
+		return nil
+	}
+
+	// Handle workloads if any exist
+	if len(groupWorkloads) > 0 {
+		if withWorkloadsFlag {
+			err = deleteWorkloadsInGroup(ctx, groupWorkloads, groupName)
+		} else {
+			err = removeWorkloadsMembershipFromGroup(ctx, groupWorkloads, groupName)
+		}
+	}
+	if err != nil {
+		return err
+	}
+
+	if err = groupManager.Delete(ctx, groupName); err != nil {
+		return fmt.Errorf("failed to delete group: %w", err)
+	}
+
+	fmt.Printf("Group '%s' deleted successfully\n", groupName)
+	return nil
+}
+
+func showWarningAndGetConfirmation(groupName string, groupWorkloads []string) (bool, error) {
+	if len(groupWorkloads) == 0 {
+		return true, nil
+	}
+
+	// Show warning and get user confirmation
+	if withWorkloadsFlag {
+		fmt.Printf("⚠️  WARNING: This will delete group '%s' and DELETE all workloads belonging to it.\n", groupName)
+	} else {
+		fmt.Printf("⚠️  WARNING: This will delete group '%s' and move all workloads to the 'default' group\n", groupName)
+	}
+
+	fmt.Printf("   The following %d workload(s) will be affected:\n", len(groupWorkloads))
+	for _, workload := range groupWorkloads {
+		if withWorkloadsFlag {
+			fmt.Printf("   - %s (will be DELETED)\n", workload)
+		} else {
+			fmt.Printf("   - %s (will be moved to the 'default' group)\n", workload)
+		}
+	}
+
+	if withWorkloadsFlag {
+		fmt.Printf("\nThis action cannot be undone. Are you sure you want to continue? [y/N]: ")
+	} else {
+		fmt.Printf("\nAre you sure you want to continue? [y/N]: ")
+	}
+
+	// Read user input
+	reader := bufio.NewReader(os.Stdin)
+	response, err := reader.ReadString('\n')
+	if err != nil {
+		return false, fmt.Errorf("failed to read user input: %w", err)
+	}
+
+	// Check if user confirmed
+	response = strings.TrimSpace(strings.ToLower(response))
+	if response != "y" && response != "yes" {
+		fmt.Println("Group deletion cancelled.")
+		return false, nil
+	}
+
+	return true, nil
+}
+
+func deleteWorkloadsInGroup(ctx context.Context, groupWorkloads []string, groupName string) error {
+	// Delete workloads
+	workloadManager, err := workloads.NewManager(ctx)
+	if err != nil {
+		return fmt.Errorf("failed to create workload manager: %w", err)
+	}
+
+	// Delete all workloads in the group
+	group, err := workloadManager.DeleteWorkloads(ctx, groupWorkloads)
+	if err != nil {
+		return fmt.Errorf("failed to delete workloads in group: %w", err)
+	}
+
+	// Wait for the deletion to complete
+	if err := group.Wait(); err != nil {
+		return fmt.Errorf("failed to delete workloads in group: %w", err)
+	}
+
+	fmt.Printf("Deleted %d workload(s) from group '%s'\n", len(groupWorkloads), groupName)
+	return nil
+}
+
+// removeWorkloadsFromGroup removes the group membership from the workloads
+// in the group.
+func removeWorkloadsMembershipFromGroup(ctx context.Context, groupWorkloads []string, groupName string) error {
+	workloadManager, err := workloads.NewManager(ctx)
+	if err != nil {
+		return fmt.Errorf("failed to create workload manager: %w", err)
+	}
+
+	// Remove group membership from all workloads
+	if err := workloadManager.MoveToDefaultGroup(ctx, groupWorkloads, groupName); err != nil {
+		return fmt.Errorf("failed to move workloads to default group: %w", err)
+	}
+
+	fmt.Printf("Removed %d workload(s) from group '%s'\n", len(groupWorkloads), groupName)
+	return nil
+}
+
 func init() {
 	groupCmd.AddCommand(groupCreateCmd)
 	groupCmd.AddCommand(groupListCmd)
+	groupCmd.AddCommand(groupRmCmd)
+
+	// Add --with-workloads flag to group rm command
+	groupRmCmd.Flags().BoolVar(&withWorkloadsFlag, "with-workloads", false,
+		"Delete all workloads in the group along with the group")
 }

cmd/thv/app/rm.go

@@ -1,27 +1,54 @@
 package app
 
 import (
+	"context"
 	"fmt"
 
 	"github.com/spf13/cobra"
 
+	"github.com/stacklok/toolhive/pkg/groups"
 	"github.com/stacklok/toolhive/pkg/workloads"
 )
 
 var rmCmd = &cobra.Command{
 	Use:               "rm [workload-name]",
 	Short:             "Remove an MCP server",
 	Long:              `Remove an MCP server managed by ToolHive.`,
-	Args:              cobra.ExactArgs(1),
+	Args:              cobra.MaximumNArgs(1),
 	RunE:              rmCmdFunc,
 	ValidArgsFunction: completeMCPServerNames,
 }
 
+func init() {
+	// TODO: Re-enable when group functionality is complete
+	// rmCmd.Flags().String("group", "", "Delete all workloads in the specified group")
+}
+
 //nolint:gocyclo // This function is complex but manageable
 func rmCmdFunc(cmd *cobra.Command, args []string) error {
 	ctx := cmd.Context()
-	// Get workload name
-	workloadName := args[0]
+
+	// Check if group flag is provided
+	groupName, _ := cmd.Flags().GetString("group")
+
+	workloadName := ""
+	if len(args) > 0 {
+		workloadName = args[0]
+	}
+
+	if workloadName != "" && groupName != "" {
+		return fmt.Errorf("workload name and group name cannot be used together")
+	}
+
+	if groupName != "" {
+		// Delete all workloads in the specified group
+		return deleteAllWorkloadsInGroup(ctx, groupName)
+	}
+
+	// Delete specific workload
+	if workloadName == "" {
+		return fmt.Errorf("workload name is required when not using --group flag")
+	}
 
 	// Create workload manager.
 	manager, err := workloads.NewManager(ctx)
@@ -43,3 +70,51 @@ func rmCmdFunc(cmd *cobra.Command, args []string) error {
 	fmt.Printf("Container %s removed successfully\n", workloadName)
 	return nil
 }
+
+func deleteAllWorkloadsInGroup(ctx context.Context, groupName string) error {
+	// Create group manager
+	groupManager, err := groups.NewManager()
+	if err != nil {
+		return fmt.Errorf("failed to create group manager: %v", err)
+	}
+
+	// Check if group exists
+	exists, err := groupManager.Exists(ctx, groupName)
+	if err != nil {
+		return fmt.Errorf("failed to check if group exists: %v", err)
+	}
+	if !exists {
+		return fmt.Errorf("group '%s' does not exist", groupName)
+	}
+
+	// Get all workloads in the group
+	groupWorkloads, err := groupManager.ListWorkloadsInGroup(ctx, groupName)
+	if err != nil {
+		return fmt.Errorf("failed to list workloads in group: %v", err)
+	}
+
+	if len(groupWorkloads) == 0 {
+		fmt.Printf("No workloads found in group '%s'\n", groupName)
+		return nil
+	}
+
+	// Create workload manager
+	workloadManager, err := workloads.NewManager(ctx)
+	if err != nil {
+		return fmt.Errorf("failed to create workload manager: %v", err)
+	}
+
+	// Delete all workloads in the group
+	group, err := workloadManager.DeleteWorkloads(ctx, groupWorkloads)
+	if err != nil {
+		return fmt.Errorf("failed to delete workloads in group: %v", err)
+	}
+
+	// Wait for the deletion to complete
+	if err := group.Wait(); err != nil {
+		return fmt.Errorf("failed to delete workloads in group: %v", err)
+	}
+
+	fmt.Printf("Successfully removed %d workload(s) from group '%s'\n", len(groupWorkloads), groupName)
+	return nil
+}

pkg/groups/group.go

@@ -8,6 +8,9 @@ import (
 	"os"
 )
 
+// DefaultGroup is the name of the default group for workloads
+const DefaultGroup = "default"
+
 // Group represents a logical grouping of MCP servers.
 type Group struct {
 	Name string `json:"name"`
@@ -44,4 +47,7 @@ type Manager interface {
 	// GetWorkloadGroup returns the group that a workload belongs to, if any.
 	// Returns nil if the workload is not in any group.
 	GetWorkloadGroup(ctx context.Context, workloadName string) (*Group, error)
+
+	// ListWorkloadsInGroup returns all workload names that belong to the specified group.
+	ListWorkloadsInGroup(ctx context.Context, groupName string) ([]string, error)
 }

pkg/groups/group_test.go

@@ -101,7 +101,7 @@ func TestManager_Create(t *testing.T) {
 			defer ctrl.Finish()
 
 			mockStore := mocks.NewMockStore(ctrl)
-			manager := &manager{store: mockStore}
+			manager := &manager{groupStore: mockStore}
 
 			// Set up mock expectations
 			tt.setupMock(mockStore)
@@ -176,7 +176,7 @@ func TestManager_Get(t *testing.T) {
 			defer ctrl.Finish()
 
 			mockStore := mocks.NewMockStore(ctrl)
-			manager := &manager{store: mockStore}
+			manager := &manager{groupStore: mockStore}
 
 			// Set up mock expectations
 			tt.setupMock(mockStore)
@@ -282,7 +282,7 @@ func TestManager_List(t *testing.T) {
 			defer ctrl.Finish()
 
 			mockStore := mocks.NewMockStore(ctrl)
-			manager := &manager{store: mockStore}
+			manager := &manager{groupStore: mockStore}
 
 			// Set up mock expectations
 			tt.setupMock(mockStore)
@@ -368,7 +368,7 @@ func TestManager_Delete(t *testing.T) {
 			defer ctrl.Finish()
 
 			mockStore := mocks.NewMockStore(ctrl)
-			manager := &manager{store: mockStore}
+			manager := &manager{groupStore: mockStore}
 
 			// Set up mock expectations
 			tt.setupMock(mockStore)
@@ -442,7 +442,7 @@ func TestManager_Exists(t *testing.T) {
 			defer ctrl.Finish()
 
 			mockStore := mocks.NewMockStore(ctrl)
-			manager := &manager{store: mockStore}
+			manager := &manager{groupStore: mockStore}
 
 			// Set up mock expectations
 			tt.setupMock(mockStore)
@@ -489,7 +489,7 @@ func TestManager_GetWorkloadGroup(t *testing.T) {
 			defer ctrl.Finish()
 
 			mockStore := mocks.NewMockStore(ctrl)
-			manager := &manager{store: mockStore}
+			manager := &manager{groupStore: mockStore}
 
 			// Call the method
 			group, err := manager.GetWorkloadGroup(context.Background(), tt.workloadName)