Move Get/List back into workloads.Manager (#1028)

In a previous PR, I decided to move the operations to Get/List workloads out of the main manager interface, and into the status tracker. After working on wiring up the status tracker, I've come to realize that a simpler approach is to keep the Get/List methods in the workload Manager, and have the workload manager call the status manager. This PR undoes the previous move. This has been lifted out of a bigger PR to keep the diff small.

Author: dmjb

cmd/thv/app/common.go

@@ -92,7 +92,7 @@ func completeMCPServerNames(cmd *cobra.Command, args []string, _ string) ([]stri
 	ctx := cmd.Context()
 
 	// Create status manager
-	manager, err := workloads.NewStatusManager(ctx)
+	manager, err := workloads.NewManager(ctx)
 	if err != nil {
 		return nil, cobra.ShellCompDirectiveError
 	}
@@ -124,7 +124,7 @@ func completeLogsArgs(cmd *cobra.Command, args []string, _ string) ([]string, co
 	ctx := cmd.Context()
 
 	// Create status manager
-	manager, err := workloads.NewStatusManager(ctx)
+	manager, err := workloads.NewManager(ctx)
 	if err != nil {
 		return []string{"prune"}, cobra.ShellCompDirectiveNoFileComp
 	}

cmd/thv/app/inspector.go

@@ -168,7 +168,7 @@ func inspectorCmdFunc(cmd *cobra.Command, args []string) error {
 
 func getServerPortAndTransport(ctx context.Context, serverName string) (int, types.TransportType, error) {
 	// Instantiate the status manager and list all workloads.
-	manager, err := workloads.NewStatusManager(ctx)
+	manager, err := workloads.NewManager(ctx)
 	if err != nil {
 		return 0, types.TransportTypeSSE, fmt.Errorf("failed to create status manager: %v", err)
 	}

cmd/thv/app/list.go

@@ -33,7 +33,7 @@ func listCmdFunc(cmd *cobra.Command, _ []string) error {
 	ctx := cmd.Context()
 
 	// Instantiate the status manager.
-	manager, err := workloads.NewStatusManager(ctx)
+	manager, err := workloads.NewManager(ctx)
 	if err != nil {
 		return fmt.Errorf("failed to create status manager: %v", err)
 	}

cmd/thv/app/logs.go

@@ -126,7 +126,7 @@ func getLogsDirectory() (string, error) {
 }
 
 func getManagedContainerNames(ctx context.Context) (map[string]bool, error) {
-	manager, err := workloads.NewStatusManager(ctx)
+	manager, err := workloads.NewManager(ctx)
 	if err != nil {
 		return nil, fmt.Errorf("failed to create status manager: %v", err)
 	}

cmd/thv/app/restart.go

@@ -44,13 +44,8 @@ func restartCmdFunc(cmd *cobra.Command, args []string) error {
 		return fmt.Errorf("failed to create workload manager: %v", err)
 	}
 
-	statusManager, err := workloads.NewStatusManager(ctx)
-	if err != nil {
-		return fmt.Errorf("failed to create status manager: %v", err)
-	}
-
 	if restartAll {
-		return restartAllContainers(ctx, statusManager, workloadManager)
+		return restartAllContainers(ctx, workloadManager)
 	}
 
 	// Restart single workload
@@ -69,13 +64,9 @@ func restartCmdFunc(cmd *cobra.Command, args []string) error {
 	return nil
 }
 
-func restartAllContainers(
-	ctx context.Context,
-	statusManager workloads.StatusManager,
-	workloadManager workloads.Manager,
-) error {
+func restartAllContainers(ctx context.Context, workloadManager workloads.Manager) error {
 	// Get all containers (including stopped ones since restart can start stopped containers)
-	allWorkloads, err := statusManager.ListWorkloads(ctx, true)
+	allWorkloads, err := workloadManager.ListWorkloads(ctx, true)
 	if err != nil {
 		return fmt.Errorf("failed to list allWorkloads: %v", err)
 	}

cmd/thv/app/stop.go

@@ -57,17 +57,12 @@ func stopCmdFunc(cmd *cobra.Command, args []string) error {
 		return fmt.Errorf("failed to create workload manager: %v", err)
 	}
 
-	statusManager, err := workloads.NewStatusManager(ctx)
-	if err != nil {
-		return fmt.Errorf("failed to create status manager: %v", err)
-	}
-
 	var group *errgroup.Group
 
 	// Check if --all flag is set
 	if stopAll {
 		// Get list of all running workloads first
-		workloadList, err := statusManager.ListWorkloads(ctx, false) // false = only running workloads
+		workloadList, err := workloadManager.ListWorkloads(ctx, false) // false = only running workloads
 		if err != nil {
 			return fmt.Errorf("failed to list workloads: %v", err)
 		}

pkg/api/server.go

@@ -132,13 +132,12 @@ func Serve(
 		return fmt.Errorf("failed to create client manager: %v", err)
 	}
 
-	statusManager := workloads.NewStatusManagerFromRuntime(rt)
 	workloadManager := workloads.NewManagerFromRuntime(rt)
 
 	routers := map[string]http.Handler{
 		"/health":               v1.HealthcheckRouter(rt),
 		"/api/v1beta/version":   v1.VersionRouter(),
-		"/api/v1beta/workloads": v1.WorkloadRouter(workloadManager, statusManager, rt, debugMode),
+		"/api/v1beta/workloads": v1.WorkloadRouter(workloadManager, rt, debugMode),
 		"/api/v1beta/registry":  v1.RegistryRouter(registryProvider),
 		"/api/v1beta/discovery": v1.DiscoveryRouter(),
 		"/api/v1beta/clients":   v1.ClientRouter(clientManager),

pkg/api/v1/workloads.go

@@ -22,7 +22,6 @@ import (
 // WorkloadRoutes defines the routes for workload management.
 type WorkloadRoutes struct {
 	workloadManager  workloads.Manager
-	statusManager    workloads.StatusManager
 	containerRuntime runtime.Runtime
 	debugMode        bool
 }
@@ -36,13 +35,11 @@ type WorkloadRoutes struct {
 // WorkloadRouter creates a new WorkloadRoutes instance.
 func WorkloadRouter(
 	workloadManager workloads.Manager,
-	statusManager workloads.StatusManager,
 	containerRuntime runtime.Runtime,
 	debugMode bool,
 ) http.Handler {
 	routes := WorkloadRoutes{
 		workloadManager:  workloadManager,
-		statusManager:    statusManager,
 		containerRuntime: containerRuntime,
 		debugMode:        debugMode,
 	}
@@ -73,7 +70,7 @@ func WorkloadRouter(
 func (s *WorkloadRoutes) listWorkloads(w http.ResponseWriter, r *http.Request) {
 	ctx := r.Context()
 	listAll := r.URL.Query().Get("all") == "true"
-	workloadList, err := s.statusManager.ListWorkloads(ctx, listAll)
+	workloadList, err := s.workloadManager.ListWorkloads(ctx, listAll)
 	if err != nil {
 		logger.Errorf("Failed to list workloads: %v", err)
 		http.Error(w, "Failed to list workloads", http.StatusInternalServerError)
@@ -102,7 +99,7 @@ func (s *WorkloadRoutes) getWorkload(w http.ResponseWriter, r *http.Request) {
 	ctx := r.Context()
 	name := chi.URLParam(r, "name")
 
-	workload, err := s.statusManager.GetWorkload(ctx, name)
+	workload, err := s.workloadManager.GetWorkload(ctx, name)
 	if err != nil {
 		if errors.Is(err, workloads.ErrWorkloadNotFound) {
 			http.Error(w, "Workload not found", http.StatusNotFound)

pkg/workloads/manager.go

@@ -33,6 +33,11 @@ import (
 // NOTE: This interface may be split up in future PRs, in particular, operations
 // which are only relevant to the CLI/API use case will be split out.
 type Manager interface {
+	// GetWorkload retrieves details of the named workload including its status.
+	GetWorkload(ctx context.Context, workloadName string) (Workload, error)
+	// ListWorkloads retrieves the states of all workloads.
+	// The `listAll` parameter determines whether to include workloads that are not running.
+	ListWorkloads(ctx context.Context, listAll bool) ([]Workload, error)
 	// DeleteWorkloads deletes the specified workloads by name.
 	// It is implemented as an asynchronous operation which returns an errgroup.Group
 	DeleteWorkloads(ctx context.Context, names []string) (*errgroup.Group, error)
@@ -91,6 +96,44 @@ func NewManagerFromRuntime(runtime rt.Runtime) Manager {
 	}
 }
 
+func (d *defaultManager) GetWorkload(ctx context.Context, workloadName string) (Workload, error) {
+	// Validate workload name to prevent path traversal attacks
+	if err := validateWorkloadName(workloadName); err != nil {
+		return Workload{}, err
+	}
+
+	container, err := d.findContainerByName(ctx, workloadName)
+	if err != nil {
+		// Note that `findContainerByName` already wraps the error with a more specific message.
+		return Workload{}, err
+	}
+
+	return WorkloadFromContainerInfo(container)
+}
+
+func (d *defaultManager) ListWorkloads(ctx context.Context, listAll bool) ([]Workload, error) {
+	// List containers
+	containers, err := d.runtime.ListWorkloads(ctx)
+	if err != nil {
+		return nil, fmt.Errorf("failed to list containers: %v", err)
+	}
+
+	// Filter containers to only show those managed by ToolHive
+	var workloads []Workload
+	for _, c := range containers {
+		// If the caller did not set `listAll` to true, only include running containers.
+		if labels.IsToolHiveContainer(c.Labels) && (isContainerRunning(&c) || listAll) {
+			workload, err := WorkloadFromContainerInfo(&c)
+			if err != nil {
+				return nil, err
+			}
+			workloads = append(workloads, workload)
+		}
+	}
+
+	return workloads, nil
+}
+
 func (d *defaultManager) StopWorkloads(ctx context.Context, names []string) (*errgroup.Group, error) {
 	// Validate all workload names to prevent path traversal attacks
 	for _, name := range names {

pkg/workloads/status.go

@@ -2,20 +2,13 @@ package workloads
 
 import (
 	"context"
-	"fmt"
 
 	ct "github.com/stacklok/toolhive/pkg/container"
 	rt "github.com/stacklok/toolhive/pkg/container/runtime"
-	"github.com/stacklok/toolhive/pkg/labels"
 )
 
 // StatusManager is an interface for fetching and retrieving workload statuses.
 type StatusManager interface {
-	// GetWorkload retrieves details of the named workload including its status.
-	GetWorkload(ctx context.Context, workloadName string) (Workload, error)
-	// ListWorkloads retrieves the states of all workloads.
-	// The `listAll` parameter determines whether to include workloads that are not running.
-	ListWorkloads(ctx context.Context, listAll bool) ([]Workload, error)
 	// SetWorkloadStatus sets the status of a workload by its name.
 	SetWorkloadStatus(ctx context.Context, workloadName string, status WorkloadStatus, contextMsg string) error
 	// DeleteWorkloadStatus removes the status of a workload by its name.
@@ -49,44 +42,6 @@ type runtimeStatusManager struct {
 	runtime rt.Runtime
 }
 
-func (r *runtimeStatusManager) GetWorkload(ctx context.Context, workloadName string) (Workload, error) {
-	// Validate workload name to prevent path traversal attacks
-	if err := validateWorkloadName(workloadName); err != nil {
-		return Workload{}, err
-	}
-
-	container, err := r.findContainerByName(ctx, workloadName)
-	if err != nil {
-		// Note that `findContainerByName` already wraps the error with a more specific message.
-		return Workload{}, err
-	}
-
-	return WorkloadFromContainerInfo(container)
-}
-
-func (r *runtimeStatusManager) ListWorkloads(ctx context.Context, listAll bool) ([]Workload, error) {
-	// List containers
-	containers, err := r.runtime.ListWorkloads(ctx)
-	if err != nil {
-		return nil, fmt.Errorf("failed to list containers: %v", err)
-	}
-
-	// Filter containers to only show those managed by ToolHive
-	var workloads []Workload
-	for _, c := range containers {
-		// If the caller did not set `listAll` to true, only include running containers.
-		if labels.IsToolHiveContainer(c.Labels) && (isContainerRunning(&c) || listAll) {
-			workload, err := WorkloadFromContainerInfo(&c)
-			if err != nil {
-				return nil, err
-			}
-			workloads = append(workloads, workload)
-		}
-	}
-
-	return workloads, nil
-}
-
 func (*runtimeStatusManager) SetWorkloadStatus(_ context.Context, _ string, _ WorkloadStatus, _ string) error {
 	// Noop
 	return nil
@@ -96,33 +51,3 @@ func (*runtimeStatusManager) DeleteWorkloadStatus(_ context.Context, _ string) e
 	// Noop
 	return nil
 }
-
-// Duplicated from the original code - need to de-dupe at some point.
-func (r *runtimeStatusManager) findContainerByName(ctx context.Context, name string) (*rt.ContainerInfo, error) {
-	// List containers to find the one with the given name
-	containers, err := r.runtime.ListWorkloads(ctx)
-	if err != nil {
-		return nil, fmt.Errorf("failed to list containers: %v", err)
-	}
-
-	// Find the container with the given name
-	for _, c := range containers {
-		// Check if the container is managed by ToolHive
-		if !labels.IsToolHiveContainer(c.Labels) {
-			continue
-		}
-
-		// Check if the container name matches
-		containerName := labels.GetContainerName(c.Labels)
-		if containerName == "" {
-			name = c.Name // Fallback to container name
-		}
-
-		// Check if the name matches (exact match or prefix match)
-		if containerName == name || c.ID == name {
-			return &c, nil
-		}
-	}
-
-	return nil, fmt.Errorf("%w: %s", ErrWorkloadNotFound, name)
-}