Update client config when workload group changes (#1442)

Author: eleftherias

cmd/thv/app/group.go

@@ -10,6 +10,9 @@ import (
 
 	"github.com/spf13/cobra"
 
+	"github.com/stacklok/toolhive/pkg/client"
+	runtime "github.com/stacklok/toolhive/pkg/container/runtime"
+	"github.com/stacklok/toolhive/pkg/core"
 	"github.com/stacklok/toolhive/pkg/groups"
 	"github.com/stacklok/toolhive/pkg/validation"
 	"github.com/stacklok/toolhive/pkg/workloads"
@@ -141,10 +144,15 @@ func groupRmCmdFunc(cmd *cobra.Command, args []string) error {
 		return fmt.Errorf("failed to create workloads manager: %w", err)
 	}
 
-	// Get all workloads in the group
-	groupWorkloads, err := workloadsManager.ListWorkloadsInGroup(ctx, groupName)
+	// Get all workloads and filter for the group
+	allWorkloads, err := workloadsManager.ListWorkloads(ctx, true) // listAll=true to include stopped workloads
 	if err != nil {
-		return fmt.Errorf("failed to list workloads in group: %w", err)
+		return fmt.Errorf("failed to list workloads: %w", err)
+	}
+
+	groupWorkloads, err := workloads.FilterByGroup(allWorkloads, groupName)
+	if err != nil {
+		return fmt.Errorf("failed to filter workloads by group: %w", err)
 	}
 
 	// Show warning and get user confirmation
@@ -160,9 +168,9 @@ func groupRmCmdFunc(cmd *cobra.Command, args []string) error {
 	// Handle workloads if any exist
 	if len(groupWorkloads) > 0 {
 		if withWorkloadsFlag {
-			err = deleteWorkloadsInGroup(ctx, groupWorkloads, groupName)
+			err = deleteWorkloadsInGroup(ctx, workloadsManager, groupWorkloads, groupName)
 		} else {
-			err = removeWorkloadsMembershipFromGroup(ctx, groupWorkloads, groupName)
+			err = moveWorkloadsToGroup(ctx, workloadsManager, groupWorkloads, groupName, groups.DefaultGroup)
 		}
 	}
 	if err != nil {
@@ -177,7 +185,7 @@ func groupRmCmdFunc(cmd *cobra.Command, args []string) error {
 	return nil
 }
 
-func showWarningAndGetConfirmation(groupName string, groupWorkloads []string) (bool, error) {
+func showWarningAndGetConfirmation(groupName string, groupWorkloads []core.Workload) (bool, error) {
 	if len(groupWorkloads) == 0 {
 		return true, nil
 	}
@@ -192,9 +200,9 @@ func showWarningAndGetConfirmation(groupName string, groupWorkloads []string) (b
 	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)
+			fmt.Printf("   - %s (will be DELETED)\n", workload.Name)
 		} else {
-			fmt.Printf("   - %s (will be moved to the 'default' group)\n", workload)
+			fmt.Printf("   - %s (will be moved to the 'default' group)\n", workload.Name)
 		}
 	}
 
@@ -221,15 +229,20 @@ func showWarningAndGetConfirmation(groupName string, groupWorkloads []string) (b
 	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)
+func deleteWorkloadsInGroup(
+	ctx context.Context,
+	workloadManager workloads.Manager,
+	groupWorkloads []core.Workload,
+	groupName string,
+) error {
+	// Extract workload names for deletion
+	var workloadNames []string
+	for _, workload := range groupWorkloads {
+		workloadNames = append(workloadNames, workload.Name)
 	}
 
 	// Delete all workloads in the group
-	group, err := workloadManager.DeleteWorkloads(ctx, groupWorkloads)
+	group, err := workloadManager.DeleteWorkloads(ctx, workloadNames)
 	if err != nil {
 		return fmt.Errorf("failed to delete workloads in group: %w", err)
 	}
@@ -243,20 +256,56 @@ func deleteWorkloadsInGroup(ctx context.Context, groupWorkloads []string, groupN
 	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)
+// moveWorkloadsToGroup moves all workloads in the specified group to a new group.
+func moveWorkloadsToGroup(
+	ctx context.Context,
+	workloadManager workloads.Manager,
+	groupWorkloads []core.Workload,
+	groupFrom string,
+	groupTo string,
+) error {
+
+	// Extract workload names for the move operation
+	var workloadNames []string
+	for _, workload := range groupWorkloads {
+		workloadNames = append(workloadNames, workload.Name)
 	}
 
-	// Remove group membership from all workloads
-	if err := workloadManager.MoveToDefaultGroup(ctx, groupWorkloads, groupName); err != nil {
+	// Update workload runconfigs to point to the new group
+	if err := workloadManager.MoveToGroup(ctx, workloadNames, groupFrom, groupTo); 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)
+	// Update client configurations for the moved workloads
+	err := updateClientConfigurations(ctx, groupWorkloads, groupFrom, groupTo)
+	if err != nil {
+		return fmt.Errorf("failed to update client configurations with new group: %w", err)
+	}
+
+	fmt.Printf("Moved %d workload(s) from group '%s' to group '%s'\n", len(groupWorkloads), groupFrom, groupTo)
+	return nil
+}
+
+func updateClientConfigurations(ctx context.Context, groupWorkloads []core.Workload, groupFrom string, groupTo string) error {
+	clientManager, err := client.NewManager(ctx)
+	if err != nil {
+		return fmt.Errorf("failed to create client manager: %w", err)
+	}
+
+	for _, w := range groupWorkloads {
+		// Only update client configurations for running workloads
+		if w.Status != runtime.WorkloadStatusRunning {
+			continue
+		}
+
+		if err := clientManager.RemoveServerFromClients(ctx, w.Name, groupFrom); err != nil {
+			return fmt.Errorf("failed to remove server %s from client configurations: %w", w.Name, err)
+		}
+		if err := clientManager.AddServerToClients(ctx, w.Name, w.URL, string(w.TransportType), groupTo); err != nil {
+			return fmt.Errorf("failed to add server %s to client configurations: %w", w.Name, err)
+		}
+	}
+
 	return nil
 }
 

pkg/api/server.go

@@ -190,7 +190,7 @@ func Serve(
 		"/api/v1beta/discovery": v1.DiscoveryRouter(),
 		"/api/v1beta/clients":   v1.ClientRouter(clientManager, workloadManager, groupManager),
 		"/api/v1beta/secrets":   v1.SecretsRouter(),
-		"/api/v1beta/groups":    v1.GroupsRouter(groupManager, workloadManager),
+		"/api/v1beta/groups":    v1.GroupsRouter(groupManager, workloadManager, clientManager),
 	}
 
 	// Only mount docs router if enabled

pkg/api/v1/groups.go

@@ -1,12 +1,16 @@
 package v1
 
 import (
+	"context"
 	"encoding/json"
 	"fmt"
 	"net/http"
 
 	"github.com/go-chi/chi/v5"
 
+	"github.com/stacklok/toolhive/pkg/client"
+	"github.com/stacklok/toolhive/pkg/container/runtime"
+	"github.com/stacklok/toolhive/pkg/core"
 	"github.com/stacklok/toolhive/pkg/errors"
 	"github.com/stacklok/toolhive/pkg/groups"
 	"github.com/stacklok/toolhive/pkg/logger"
@@ -18,13 +22,15 @@ import (
 type GroupsRoutes struct {
 	groupManager    groups.Manager
 	workloadManager workloads.Manager
+	clientManager   client.Manager
 }
 
 // GroupsRouter creates a new GroupsRoutes instance.
-func GroupsRouter(groupManager groups.Manager, workloadManager workloads.Manager) http.Handler {
+func GroupsRouter(groupManager groups.Manager, workloadManager workloads.Manager, clientManager client.Manager) http.Handler {
 	routes := GroupsRoutes{
 		groupManager:    groupManager,
 		workloadManager: workloadManager,
+		clientManager:   clientManager,
 	}
 
 	r := chi.NewRouter()
@@ -202,42 +208,27 @@ func (s *GroupsRoutes) deleteGroup(w http.ResponseWriter, r *http.Request) {
 	// Get the with-workloads flag from query parameter
 	withWorkloads := r.URL.Query().Get("with-workloads") == "true"
 
-	// Get all workloads in the group
-	groupWorkloads, err := s.workloadManager.ListWorkloadsInGroup(ctx, name)
+	// Get all workloads and filter for the group
+	allWorkloads, err := s.workloadManager.ListWorkloads(ctx, true) // listAll=true to include stopped workloads
 	if err != nil {
-		logger.Errorf("Failed to list workloads in group %s: %v", name, err)
-		http.Error(w, "Failed to list workloads in group", http.StatusInternalServerError)
+		logger.Errorf("Failed to list workloads: %v", err)
+		http.Error(w, "Failed to list workloads", http.StatusInternalServerError)
+		return
+	}
+
+	groupWorkloads, err := workloads.FilterByGroup(allWorkloads, name)
+	if err != nil {
+		logger.Errorf("Failed to filter workloads by group %s: %v", name, err)
+		http.Error(w, "Failed to filter workloads by group", http.StatusInternalServerError)
 		return
 	}
 
 	// Handle workloads if any exist
 	if len(groupWorkloads) > 0 {
-		if withWorkloads {
-			// Delete all workloads in the group
-			group, err := s.workloadManager.DeleteWorkloads(ctx, groupWorkloads)
-			if err != nil {
-				logger.Errorf("Failed to delete workloads in group %s: %v", name, err)
-				http.Error(w, "Failed to delete workloads in group", http.StatusInternalServerError)
-				return
-			}
-
-			// Wait for the deletion to complete
-			if err := group.Wait(); err != nil {
-				logger.Errorf("Failed to delete workloads in group %s: %v", name, err)
-				http.Error(w, "Failed to delete workloads in group", http.StatusInternalServerError)
-				return
-			}
-
-			logger.Infof("Deleted %d workload(s) from group '%s'", len(groupWorkloads), name)
-		} else {
-			// Move workloads to default group
-			if err := s.workloadManager.MoveToDefaultGroup(ctx, groupWorkloads, name); err != nil {
-				logger.Errorf("Failed to move workloads to default group: %v", err)
-				http.Error(w, "Failed to move workloads to default group", http.StatusInternalServerError)
-				return
-			}
-
-			logger.Infof("Moved %d workload(s) from group '%s' to default group", len(groupWorkloads), name)
+		if err := s.handleWorkloadsForGroupDeletion(ctx, name, groupWorkloads, withWorkloads); err != nil {
+			logger.Errorf("Failed to handle workloads for group %s: %v", name, err)
+			http.Error(w, "Failed to handle workloads", http.StatusInternalServerError)
+			return
 		}
 	}
 
@@ -252,6 +243,73 @@ func (s *GroupsRoutes) deleteGroup(w http.ResponseWriter, r *http.Request) {
 	w.WriteHeader(http.StatusNoContent)
 }
 
+// handleWorkloadsForGroupDeletion handles workloads when deleting a group
+func (s *GroupsRoutes) handleWorkloadsForGroupDeletion(
+	ctx context.Context,
+	groupName string,
+	groupWorkloads []core.Workload,
+	withWorkloads bool,
+) error {
+	// Extract workload names
+	var workloadNames []string
+	for _, workload := range groupWorkloads {
+		workloadNames = append(workloadNames, workload.Name)
+	}
+
+	if withWorkloads {
+		// Delete all workloads in the group
+		group, err := s.workloadManager.DeleteWorkloads(ctx, workloadNames)
+		if err != nil {
+			return fmt.Errorf("failed to delete workloads in group %s: %w", groupName, err)
+		}
+
+		// Wait for the deletion to complete
+		if err := group.Wait(); err != nil {
+			return fmt.Errorf("failed to delete workloads in group %s: %w", groupName, err)
+		}
+
+		logger.Infof("Deleted %d workload(s) from group '%s'", len(groupWorkloads), groupName)
+	} else {
+		// Move workloads to default group
+		if err := s.workloadManager.MoveToGroup(ctx, workloadNames, groupName, groups.DefaultGroup); err != nil {
+			return fmt.Errorf("failed to move workloads to default group: %w", err)
+		}
+
+		// Update client configurations for the moved workloads
+		if err := s.updateClientConfigurations(ctx, groupWorkloads, groupName, groups.DefaultGroup); err != nil {
+			return fmt.Errorf("failed to update client configurations: %w", err)
+		}
+
+		logger.Infof("Moved %d workload(s) from group '%s' to default group", len(groupWorkloads), groupName)
+	}
+
+	return nil
+}
+
+// updateClientConfigurations updates client configurations when workloads are moved between groups
+func (s *GroupsRoutes) updateClientConfigurations(
+	ctx context.Context,
+	groupWorkloads []core.Workload,
+	groupFrom string,
+	groupTo string,
+) error {
+	for _, w := range groupWorkloads {
+		// Only update client configurations for running workloads
+		if w.Status != runtime.WorkloadStatusRunning {
+			continue
+		}
+
+		if err := s.clientManager.RemoveServerFromClients(ctx, w.Name, groupFrom); err != nil {
+			return fmt.Errorf("failed to remove server %s from client configurations: %w", w.Name, err)
+		}
+		if err := s.clientManager.AddServerToClients(ctx, w.Name, w.URL, string(w.TransportType), groupTo); err != nil {
+			return fmt.Errorf("failed to add server %s to client configurations: %w", w.Name, err)
+		}
+	}
+
+	return nil
+}
+
 // Response types
 
 type groupListResponse struct {