Migrate individual servers to default group (#1186)

Author: amirejaz

cmd/thv/app/run.go

@@ -127,36 +127,9 @@ func runCmdFunc(cmd *cobra.Command, args []string) error {
 	// Get debug mode flag
 	debugMode, _ := cmd.Flags().GetBool("debug")
 
-	workloadName := runFlags.Name
-	if workloadName == "" {
-		workloadName = serverOrImage
-	}
-
-	isForeground := runFlags.Foreground
-
-	if runFlags.Group != "" {
-		groupManager, err := groups.NewManager()
-		if err != nil {
-			return fmt.Errorf("failed to create group manager: %v", err)
-		}
-
-		// Check if the workload is already in a group
-		group, err := groupManager.GetWorkloadGroup(ctx, workloadName)
-		if err != nil {
-			return fmt.Errorf("failed to get workload group: %v", err)
-		}
-		if group != nil && group.Name != runFlags.Group {
-			return fmt.Errorf("workload '%s' is already in group '%s'", workloadName, group.Name)
-		}
-
-		// Validate that the group specified exists
-		exists, err := groupManager.Exists(ctx, runFlags.Group)
-		if err != nil {
-			return fmt.Errorf("failed to check if group exists: %v", err)
-		}
-		if !exists {
-			return fmt.Errorf("group '%s' does not exist", runFlags.Group)
-		}
+	err := validateGroup(ctx, serverOrImage)
+	if err != nil {
+		return err
 	}
 
 	// Build the run configuration
@@ -172,7 +145,7 @@ func runCmdFunc(cmd *cobra.Command, args []string) error {
 	}
 	workloadManager := workloads.NewManagerFromRuntime(rt)
 
-	if isForeground {
+	if runFlags.Foreground {
 		return runForeground(ctx, workloadManager, runnerConfig)
 	}
 	return workloadManager.RunWorkloadDetached(ctx, runnerConfig)
@@ -203,6 +176,40 @@ func runForeground(ctx context.Context, workloadManager workloads.Manager, runne
 	}
 }
 
+func validateGroup(ctx context.Context, serverOrImage string) error {
+	workloadName := runFlags.Name
+	if workloadName == "" {
+		workloadName = serverOrImage
+	}
+
+	// Create group manager
+	groupManager, err := groups.NewManager()
+	if err != nil {
+		return fmt.Errorf("failed to create group manager: %v", err)
+	}
+
+	// Check if the workload is already in a group
+	group, err := groupManager.GetWorkloadGroup(ctx, workloadName)
+	if err != nil {
+		return fmt.Errorf("failed to get workload group: %v", err)
+	}
+	if group != nil && group.Name != runFlags.Group {
+		return fmt.Errorf("workload '%s' is already in group '%s'", workloadName, group.Name)
+	}
+
+	if runFlags.Group != "" {
+		// Validate that the group specified exists
+		exists, err := groupManager.Exists(ctx, runFlags.Group)
+		if err != nil {
+			return fmt.Errorf("failed to check if group exists: %v", err)
+		}
+		if !exists {
+			return fmt.Errorf("group '%s' does not exist", runFlags.Group)
+		}
+	}
+	return nil
+}
+
 // parseCommandArguments processes command-line arguments to find everything after the -- separator
 // which are the arguments to be passed to the MCP server
 func parseCommandArguments(args []string) []string {

cmd/thv/app/run_flags.go

@@ -81,7 +81,8 @@ func AddRunFlags(cmd *cobra.Command, config *RunFlags) {
 	cmd.Flags().StringVar(&config.ProxyMode, "proxy-mode", "sse", "Proxy mode for stdio transport (sse or streamable-http)")
 	cmd.Flags().StringVar(&config.Name, "name", "", "Name of the MCP server (auto-generated from image if not provided)")
 	// TODO: Re-enable when group functionality is complete
-	// cmd.Flags().StringVar(&config.Group, "group", "", "Name of the group this workload belongs to")
+	// cmd.Flags().StringVar(&config.Group, "group", "default",
+	//	"Name of the group this workload belongs to (defaults to 'default' if not specified)")
 	cmd.Flags().StringVar(&config.Host, "host", transport.LocalhostIPv4, "Host for the HTTP proxy to listen on (IP or hostname)")
 	cmd.Flags().IntVar(&config.ProxyPort, "proxy-port", 0, "Port for the HTTP proxy to listen on (host port)")
 	cmd.Flags().IntVar(&config.TargetPort, "target-port", 0,

cmd/thv/main.go

@@ -18,6 +18,11 @@ func main() {
 	// Handles the auto-discovery flag depreciation, only executes once on old config files
 	client.CheckAndPerformAutoDiscoveryMigration()
 
+	// Check and perform default group migration if needed
+	// Migrates existing workloads to the default group, only executes once
+	// TODO: Re-enable when group functionality is complete
+	// groups.CheckAndPerformDefaultGroupMigration()
+
 	// Skip update check for completion command or if we are running in kubernetes
 	if err := app.NewRootCmd(!app.IsCompletionCommand(os.Args) && !runtime.IsKubernetesRuntime()).Execute(); err != nil {
 		os.Exit(1)

pkg/config/config.go

@@ -29,6 +29,7 @@ type Config struct {
 	AllowPrivateRegistryIp bool                `yaml:"allow_private_registry_ip"`
 	CACertificatePath      string              `yaml:"ca_certificate_path,omitempty"`
 	OTEL                   OpenTelemetryConfig `yaml:"otel,omitempty"`
+	DefaultGroupMigration  bool                `yaml:"default_group_migration,omitempty"`
 }
 
 // Secrets contains the settings for secrets management.
@@ -93,6 +94,7 @@ func createNewConfigWithDefaults() Config {
 		},
 		RegistryUrl:            "",
 		AllowPrivateRegistryIp: false,
+		DefaultGroupMigration:  false,
 	}
 }
 

pkg/groups/manager.go

@@ -11,6 +11,11 @@ import (
 	"github.com/stacklok/toolhive/pkg/state"
 )
 
+const (
+	// DefaultGroupName is the name of the default group
+	DefaultGroupName = "default"
+)
+
 // manager implements the Manager interface
 type manager struct {
 	groupStore     state.Store

pkg/groups/migration.go

@@ -0,0 +1,109 @@
+package groups
+
+import (
+	"context"
+	"fmt"
+	"sync"
+
+	"github.com/stacklok/toolhive/pkg/config"
+	"github.com/stacklok/toolhive/pkg/logger"
+	"github.com/stacklok/toolhive/pkg/runner"
+	"github.com/stacklok/toolhive/pkg/state"
+)
+
+// migrationOnce ensures the migration only runs once
+var migrationOnce sync.Once
+
+// CheckAndPerformDefaultGroupMigration checks if default group migration is needed and performs it
+// This is called once at application startup
+func CheckAndPerformDefaultGroupMigration() {
+	migrationOnce.Do(func() {
+		appConfig := config.GetConfig()
+
+		// Check if default group migration has already been performed
+		if appConfig.DefaultGroupMigration {
+			return
+		}
+
+		performDefaultGroupMigration()
+	})
+}
+
+// performDefaultGroupMigration migrates all existing workloads to the default group
+func performDefaultGroupMigration() {
+	fmt.Println("Migrating existing workloads to default group...")
+	fmt.Println()
+
+	// Create group manager and ensure default group exists
+	groupManager, err := NewManager()
+	if err != nil {
+		logger.Errorf("Failed to create group manager: %v", err)
+		return
+	}
+
+	// Create default group
+	if err := createDefaultGroup(context.Background(), groupManager); err != nil {
+		logger.Errorf("Failed to create default group: %v", err)
+		return
+	}
+
+	// Create a runconfig store to list all runconfigs
+	runConfigStore, err := state.NewRunConfigStore("toolhive")
+	if err != nil {
+		logger.Errorf("Failed to create runconfig store: %v", err)
+		return
+	}
+
+	// List all runconfig names
+	runConfigNames, err := runConfigStore.List(context.Background())
+	if err != nil {
+		logger.Errorf("Failed to list runconfigs: %v", err)
+		return
+	}
+
+	migratedCount := 0
+	for _, runConfigName := range runConfigNames {
+		// Load the runconfig
+		runnerInstance, err := runner.LoadState(context.Background(), runConfigName)
+		if err != nil {
+			// Log the error but continue processing other runconfigs
+			logger.Warnf("Failed to load runconfig %s: %v", runConfigName, err)
+			continue
+		}
+
+		// If the workload has no group, assign it to the default group
+		if runnerInstance.Config.Group == "" {
+			runnerInstance.Config.Group = DefaultGroupName
+			if err := runnerInstance.SaveState(context.Background()); err != nil {
+				logger.Warnf("Failed to save runconfig for %s: %v", runConfigName, err)
+				continue
+			}
+			migratedCount++
+		}
+	}
+
+	if migratedCount > 0 {
+		fmt.Printf("\nSuccessfully migrated %d workloads to default group '%s'\n", migratedCount, DefaultGroupName)
+	} else {
+		fmt.Println("No workloads needed migration to default group")
+	}
+
+	// Mark default group migration as completed
+	err = config.UpdateConfig(func(c *config.Config) {
+		c.DefaultGroupMigration = true
+	})
+
+	if err != nil {
+		logger.Errorf("Error updating config during migration: %v", err)
+		return
+	}
+}
+
+// createDefaultGroup creates the default group if it doesn't exist
+func createDefaultGroup(ctx context.Context, groupManager Manager) error {
+	logger.Infof("Creating default group '%s'", DefaultGroupName)
+	if err := groupManager.Create(ctx, DefaultGroupName); err != nil {
+		return fmt.Errorf("failed to create default group: %w", err)
+	}
+	return nil
+}