Add group flag to client removal (#1371)

Author: eleftherias

cmd/thv/app/client.go

@@ -12,6 +12,7 @@ import (
 	"github.com/stacklok/toolhive/pkg/config"
 	"github.com/stacklok/toolhive/pkg/core"
 	"github.com/stacklok/toolhive/pkg/groups"
+	"github.com/stacklok/toolhive/pkg/logger"
 	"github.com/stacklok/toolhive/pkg/workloads"
 )
 
@@ -104,6 +105,8 @@ func init() {
 	// TODO: Re-enable when group functionality is complete
 	//clientRegisterCmd.Flags().StringSliceVar(
 	//	&groupNames, "group", []string{"default"}, "Only register workloads from specified groups")
+	//clientRemoveCmd.Flags().StringSliceVar(
+	//	&groupNames, "group", []string{}, "Remove client from specified groups (if not set, removes all workloads from the client)")
 }
 
 func clientStatusCmdFunc(_ *cobra.Command, _ []string) error {
@@ -208,21 +211,7 @@ func clientRemoveCmdFunc(cmd *cobra.Command, args []string) error {
 			clientType)
 	}
 
-	ctx := cmd.Context()
-
-	manager, err := client.NewManager(ctx)
-	if err != nil {
-		return fmt.Errorf("failed to create client manager: %w", err)
-	}
-
-	err = manager.UnregisterClients(ctx, []client.Client{
-		{Name: client.MCPClient(clientType)},
-	})
-	if err != nil {
-		return fmt.Errorf("failed to remove client %s: %w", clientType, err)
-	}
-
-	return nil
+	return performClientRemoval(cmd.Context(), client.Client{Name: client.MCPClient(clientType)}, groupNames)
 }
 
 func listRegisteredClientsCmdFunc(_ *cobra.Command, _ []string) error {
@@ -339,3 +328,114 @@ func registerClientsGlobally(
 
 	return nil
 }
+
+func performClientRemoval(ctx context.Context, clientToRemove client.Client, groupNames []string) error {
+	clientManager, err := client.NewManager(ctx)
+	if err != nil {
+		return fmt.Errorf("failed to create client manager: %w", err)
+	}
+
+	workloadManager, err := workloads.NewManager(ctx)
+	if err != nil {
+		return fmt.Errorf("failed to create workload manager: %w", err)
+	}
+
+	runningWorkloads, err := workloadManager.ListWorkloads(ctx, false)
+	if err != nil {
+		return fmt.Errorf("failed to list running workloads: %w", err)
+	}
+
+	groupManager, err := groups.NewManager()
+	if err != nil {
+		return fmt.Errorf("failed to create group manager: %w", err)
+	}
+
+	if len(groupNames) > 0 {
+		return removeClientFromGroups(ctx, clientToRemove, groupNames, runningWorkloads, groupManager, clientManager)
+	}
+
+	return removeClientGlobally(ctx, clientToRemove, runningWorkloads, groupManager, clientManager)
+}
+
+func removeClientFromGroups(
+	ctx context.Context,
+	clientToRemove client.Client,
+	groupNames []string,
+	runningWorkloads []core.Workload,
+	groupManager groups.Manager,
+	clientManager client.Manager,
+) error {
+	fmt.Printf("Filtering workloads to groups: %v\n", groupNames)
+
+	// Remove client from specific groups only
+	filteredWorkloads, err := workloads.FilterByGroups(runningWorkloads, groupNames)
+	if err != nil {
+		return fmt.Errorf("failed to filter workloads by groups: %w", err)
+	}
+
+	// Remove the workloads from the client's configuration file
+	err = clientManager.UnregisterClients(ctx, []client.Client{clientToRemove}, filteredWorkloads)
+	if err != nil {
+		return fmt.Errorf("failed to unregister client: %w", err)
+	}
+
+	// Remove the client from the groups
+	err = groupManager.UnregisterClients(ctx, groupNames, []string{string(clientToRemove.Name)})
+	if err != nil {
+		return fmt.Errorf("failed to unregister client from groups: %w", err)
+	}
+
+	fmt.Printf("Successfully removed client %s from groups: %v\n", clientToRemove.Name, groupNames)
+
+	return nil
+}
+
+func removeClientGlobally(
+	ctx context.Context,
+	clientToRemove client.Client,
+	runningWorkloads []core.Workload,
+	groupManager groups.Manager,
+	clientManager client.Manager,
+) error {
+	// Remove the workloads from the client's configuration file
+	err := clientManager.UnregisterClients(ctx, []client.Client{clientToRemove}, runningWorkloads)
+	if err != nil {
+		return fmt.Errorf("failed to unregister client: %w", err)
+	}
+
+	allGroups, err := groupManager.List(ctx)
+	if err != nil {
+		return fmt.Errorf("failed to list groups: %w", err)
+	}
+
+	if len(allGroups) > 0 {
+		// Remove client from all groups first
+		allGroupNames := make([]string, len(allGroups))
+		for i, group := range allGroups {
+			allGroupNames[i] = group.Name
+		}
+
+		err = groupManager.UnregisterClients(ctx, allGroupNames, []string{string(clientToRemove.Name)})
+		if err != nil {
+			return fmt.Errorf("failed to unregister client from groups: %w", err)
+		}
+	}
+
+	// Remove client from global registered clients list
+	err = config.UpdateConfig(func(c *config.Config) {
+		for i, registeredClient := range c.Clients.RegisteredClients {
+			if registeredClient == string(clientToRemove.Name) {
+				// Remove client from slice
+				c.Clients.RegisteredClients = append(c.Clients.RegisteredClients[:i], c.Clients.RegisteredClients[i+1:]...)
+				logger.Infof("Successfully unregistered client: %s\n", clientToRemove.Name)
+				return
+			}
+		}
+		logger.Warnf("Client %s was not found in registered clients list", clientToRemove.Name)
+	})
+	if err != nil {
+		return fmt.Errorf("failed to update configuration for client %s: %w", clientToRemove.Name, err)
+	}
+
+	return nil
+}

docs/server/docs.go

@@ -10,6 +10,7 @@ import (
 
 	"github.com/stacklok/toolhive/pkg/client"
 	"github.com/stacklok/toolhive/pkg/config"
+	"github.com/stacklok/toolhive/pkg/core"
 	"github.com/stacklok/toolhive/pkg/groups"
 	"github.com/stacklok/toolhive/pkg/logger"
 	"github.com/stacklok/toolhive/pkg/workloads"
@@ -38,6 +39,7 @@ func ClientRouter(
 	r.Get("/", routes.listClients)
 	r.Post("/", routes.registerClient)
 	r.Delete("/{name}", routes.unregisterClient)
+	r.Delete("/{name}/groups/{group}", routes.unregisterClientFromGroup)
 	r.Post("/register", routes.registerClientsBulk)
 	r.Post("/unregister", routes.unregisterClientsBulk)
 	return r
@@ -129,9 +131,7 @@ func (c *ClientRoutes) unregisterClient(w http.ResponseWriter, r *http.Request)
 		return
 	}
 
-	err := c.clientManager.UnregisterClients(r.Context(), []client.Client{
-		{Name: client.MCPClient(clientName)},
-	})
+	err := c.removeClient(r.Context(), []client.Client{{Name: client.MCPClient(clientName)}}, nil)
 	if err != nil {
 		logger.Errorf("Failed to unregister client: %v", err)
 		http.Error(w, "Failed to unregister client", http.StatusInternalServerError)
@@ -141,6 +141,41 @@ func (c *ClientRoutes) unregisterClient(w http.ResponseWriter, r *http.Request)
 	w.WriteHeader(http.StatusNoContent)
 }
 
+// unregisterClientFromGroup
+//
+//	@Summary		Unregister a client from a specific group
+//	@Description	Unregister a client from a specific group in ToolHive
+//	@Tags			clients
+//	@Param			name	path	string	true	"Client name to unregister"
+//	@Param			group	path	string	true	"Group name to remove client from"
+//	@Success		204
+//	@Failure		400	{string}	string	"Invalid request"
+//	@Failure		404	{string}	string	"Client or group not found"
+//	@Router			/api/v1beta/clients/{name}/groups/{group} [delete]
+func (c *ClientRoutes) unregisterClientFromGroup(w http.ResponseWriter, r *http.Request) {
+	clientName := chi.URLParam(r, "name")
+	if clientName == "" {
+		http.Error(w, "Client name is required", http.StatusBadRequest)
+		return
+	}
+
+	groupName := chi.URLParam(r, "group")
+	if groupName == "" {
+		http.Error(w, "Group name is required", http.StatusBadRequest)
+		return
+	}
+
+	// Remove client from the specific group
+	err := c.removeClient(r.Context(), []client.Client{{Name: client.MCPClient(clientName)}}, []string{groupName})
+	if err != nil {
+		logger.Errorf("Failed to unregister client from group: %v", err)
+		http.Error(w, "Failed to unregister client from group", http.StatusInternalServerError)
+		return
+	}
+
+	w.WriteHeader(http.StatusNoContent)
+}
+
 // registerClientsBulk
 //
 //	@Summary		Register multiple clients
@@ -220,7 +255,7 @@ func (c *ClientRoutes) unregisterClientsBulk(w http.ResponseWriter, r *http.Requ
 		clients[i] = client.Client{Name: name}
 	}
 
-	err = c.clientManager.UnregisterClients(r.Context(), clients)
+	err = c.removeClient(r.Context(), clients, req.Groups)
 	if err != nil {
 		logger.Errorf("Failed to unregister clients: %v", err)
 		http.Error(w, "Failed to unregister clients", http.StatusInternalServerError)
@@ -310,3 +345,104 @@ func (c *ClientRoutes) performClientRegistration(ctx context.Context, clients []
 
 	return nil
 }
+
+func (c *ClientRoutes) removeClient(ctx context.Context, clients []client.Client, groupNames []string) error {
+	runningWorkloads, err := c.workloadManager.ListWorkloads(ctx, false)
+	if err != nil {
+		return fmt.Errorf("failed to list running workloads: %w", err)
+	}
+
+	if len(groupNames) > 0 {
+		return c.removeClientFromGroup(ctx, clients, groupNames, runningWorkloads)
+	}
+
+	return c.removeClientGlobally(ctx, clients, runningWorkloads)
+}
+
+func (c *ClientRoutes) removeClientFromGroup(
+	ctx context.Context,
+	clients []client.Client,
+	groupNames []string,
+	runningWorkloads []core.Workload,
+) error {
+	// Remove clients from specific groups only
+	filteredWorkloads, err := workloads.FilterByGroups(runningWorkloads, groupNames)
+	if err != nil {
+		return fmt.Errorf("failed to filter workloads by groups: %w", err)
+	}
+
+	// Remove the workloads from the client's configuration file
+	err = c.clientManager.UnregisterClients(ctx, clients, filteredWorkloads)
+	if err != nil {
+		return fmt.Errorf("failed to unregister clients: %w", err)
+	}
+
+	// Extract client names for group management
+	clientNames := make([]string, len(clients))
+	for i, clientToRemove := range clients {
+		clientNames[i] = string(clientToRemove.Name)
+	}
+
+	// Remove the clients from the groups
+	err = c.groupManager.UnregisterClients(ctx, groupNames, clientNames)
+	if err != nil {
+		return fmt.Errorf("failed to unregister clients from groups: %w", err)
+	}
+
+	return nil
+}
+
+func (c *ClientRoutes) removeClientGlobally(
+	ctx context.Context,
+	clients []client.Client,
+	runningWorkloads []core.Workload,
+) error {
+	// Remove the workloads from the client's configuration file
+	err := c.clientManager.UnregisterClients(ctx, clients, runningWorkloads)
+	if err != nil {
+		return fmt.Errorf("failed to unregister clients: %w", err)
+	}
+
+	// Remove clients from all groups and global registry
+	allGroups, err := c.groupManager.List(ctx)
+	if err != nil {
+		return fmt.Errorf("failed to list groups: %w", err)
+	}
+
+	if len(allGroups) > 0 {
+		clientNames := make([]string, len(clients))
+		for i, clientToRemove := range clients {
+			clientNames[i] = string(clientToRemove.Name)
+		}
+
+		allGroupNames := make([]string, len(allGroups))
+		for i, group := range allGroups {
+			allGroupNames[i] = group.Name
+		}
+
+		err = c.groupManager.UnregisterClients(ctx, allGroupNames, clientNames)
+		if err != nil {
+			return fmt.Errorf("failed to unregister clients from groups: %w", err)
+		}
+	}
+
+	// Remove clients from global registered clients list
+	for _, clientToRemove := range clients {
+		err := config.UpdateConfig(func(c *config.Config) {
+			for i, registeredClient := range c.Clients.RegisteredClients {
+				if registeredClient == string(clientToRemove.Name) {
+					// Remove client from slice
+					c.Clients.RegisteredClients = append(c.Clients.RegisteredClients[:i], c.Clients.RegisteredClients[i+1:]...)
+					logger.Infof("Successfully unregistered client: %s\n", clientToRemove.Name)
+					return
+				}
+			}
+			logger.Warnf("Client %s was not found in registered clients list", clientToRemove.Name)
+		})
+		if err != nil {
+			return fmt.Errorf("failed to update configuration for client %s: %w", clientToRemove.Name, err)
+		}
+	}
+
+	return nil
+}