Fix bugs found while testing file-based state tracking (#1347)

This PR was intended to enable file-based state tracking. However, I found some bugs during testing which involves touching many files, so this PR addresses these bugs without enabling file-based state tracking.

1. I discovered that I left out a call to change the state to "running" once the proxy-runner is ready. Simply adding the state manager type led to a circular dependency, which required a bunch of refactoring of code. This is responsible for most of the diff here.
2. The code which created the base directory for the status files was using the wrong path - this was fixed.
3. I originally had separate methods for creating the initial status, and updating. This led to problems with the detached process flow (which would try to create the status twice, leading to an error). This PR changes the code to use a single method which creates or update.

Author: dmjb

cmd/thv-proxyrunner/app/run.go

@@ -266,7 +266,10 @@ func runCmdFunc(cmd *cobra.Command, args []string) error {
 		return fmt.Errorf("failed to create RunConfig: %v", err)
 	}
 
-	workloadManager := workloads.NewManagerFromRuntime(rt)
+	workloadManager, err := workloads.NewManagerFromRuntime(rt)
+	if err != nil {
+		return fmt.Errorf("failed to create workload manager: %v", err)
+	}
 	return workloadManager.RunWorkload(ctx, runConfig)
 }
 

cmd/thv/app/run.go

@@ -144,7 +144,10 @@ func runCmdFunc(cmd *cobra.Command, args []string) error {
 	if err != nil {
 		return fmt.Errorf("failed to create container runtime: %v", err)
 	}
-	workloadManager := workloads.NewManagerFromRuntime(rt)
+	workloadManager, err := workloads.NewManagerFromRuntime(rt)
+	if err != nil {
+		return fmt.Errorf("failed to create workload manager: %v", err)
+	}
 
 	err = validateGroup(ctx, workloadManager, serverOrImage)
 	if err != nil {
@@ -298,7 +301,10 @@ func runFromConfigFile(ctx context.Context) error {
 	runConfig.Deployer = rt
 
 	// Create workload manager
-	workloadManager := workloads.NewManagerFromRuntime(rt)
+	workloadManager, err := workloads.NewManagerFromRuntime(rt)
+	if err != nil {
+		return fmt.Errorf("failed to create workload manager: %v", err)
+	}
 
 	// Run the workload based on foreground flag
 	if runFlags.Foreground {

cmd/thv/app/stop.go

@@ -9,6 +9,7 @@ import (
 
 	rt "github.com/stacklok/toolhive/pkg/container/runtime"
 	"github.com/stacklok/toolhive/pkg/workloads"
+	"github.com/stacklok/toolhive/pkg/workloads/types"
 )
 
 var stopCmd = &cobra.Command{
@@ -100,7 +101,7 @@ func stopCmdFunc(cmd *cobra.Command, args []string) error {
 			// If the workload is not found or not running, treat as a non-fatal error.
 			if errors.Is(err, rt.ErrWorkloadNotFound) ||
 				errors.Is(err, workloads.ErrWorkloadNotRunning) ||
-				errors.Is(err, workloads.ErrInvalidWorkloadName) {
+				errors.Is(err, types.ErrInvalidWorkloadName) {
 				fmt.Printf("workload %s is not running\n", workloadName)
 				return nil
 			}

pkg/api/server.go

@@ -171,7 +171,10 @@ func Serve(
 		return fmt.Errorf("failed to create client manager: %v", err)
 	}
 
-	workloadManager := workloads.NewManagerFromRuntime(containerRuntime)
+	workloadManager, err := workloads.NewManagerFromRuntime(containerRuntime)
+	if err != nil {
+		return fmt.Errorf("failed to create workload manager: %v", err)
+	}
 
 	// Create group manager
 	groupManager, err := groups.NewManager()

pkg/api/v1/workloads.go

@@ -22,6 +22,7 @@ import (
 	"github.com/stacklok/toolhive/pkg/transport/types"
 	"github.com/stacklok/toolhive/pkg/validation"
 	"github.com/stacklok/toolhive/pkg/workloads"
+	wt "github.com/stacklok/toolhive/pkg/workloads/types"
 )
 
 // WorkloadRoutes defines the routes for workload management.
@@ -135,7 +136,7 @@ func (s *WorkloadRoutes) getWorkload(w http.ResponseWriter, r *http.Request) {
 		if errors.Is(err, runtime.ErrWorkloadNotFound) {
 			http.Error(w, "Workload not found", http.StatusNotFound)
 			return
-		} else if errors.Is(err, workloads.ErrInvalidWorkloadName) {
+		} else if errors.Is(err, wt.ErrInvalidWorkloadName) {
 			http.Error(w, "Invalid workload name: "+err.Error(), http.StatusBadRequest)
 			return
 		}
@@ -169,7 +170,7 @@ func (s *WorkloadRoutes) stopWorkload(w http.ResponseWriter, r *http.Request) {
 	// Use the bulk method with a single workload
 	_, err := s.workloadManager.StopWorkloads(ctx, []string{name})
 	if err != nil {
-		if errors.Is(err, workloads.ErrInvalidWorkloadName) {
+		if errors.Is(err, wt.ErrInvalidWorkloadName) {
 			http.Error(w, "Invalid workload name: "+err.Error(), http.StatusBadRequest)
 			return
 		}
@@ -198,7 +199,7 @@ func (s *WorkloadRoutes) restartWorkload(w http.ResponseWriter, r *http.Request)
 	// Note: In the API, we always assume that the restart is a background operation
 	_, err := s.workloadManager.RestartWorkloads(ctx, []string{name}, false)
 	if err != nil {
-		if errors.Is(err, workloads.ErrInvalidWorkloadName) {
+		if errors.Is(err, wt.ErrInvalidWorkloadName) {
 			http.Error(w, "Invalid workload name: "+err.Error(), http.StatusBadRequest)
 			return
 		}
@@ -226,7 +227,7 @@ func (s *WorkloadRoutes) deleteWorkload(w http.ResponseWriter, r *http.Request)
 	// Use the bulk method with a single workload
 	_, err := s.workloadManager.DeleteWorkloads(ctx, []string{name})
 	if err != nil {
-		if errors.Is(err, workloads.ErrInvalidWorkloadName) {
+		if errors.Is(err, wt.ErrInvalidWorkloadName) {
 			http.Error(w, "Invalid workload name: "+err.Error(), http.StatusBadRequest)
 			return
 		}
@@ -366,7 +367,7 @@ func (s *WorkloadRoutes) stopWorkloadsBulk(w http.ResponseWriter, r *http.Reques
 	// The request is not blocked on completion.
 	_, err = s.workloadManager.StopWorkloads(ctx, workloadNames)
 	if err != nil {
-		if errors.Is(err, workloads.ErrInvalidWorkloadName) {
+		if errors.Is(err, wt.ErrInvalidWorkloadName) {
 			http.Error(w, "Invalid workload name: "+err.Error(), http.StatusBadRequest)
 			return
 		}
@@ -412,7 +413,7 @@ func (s *WorkloadRoutes) restartWorkloadsBulk(w http.ResponseWriter, r *http.Req
 	// Note: In the API, we always assume that the restart is a background operation.
 	_, err = s.workloadManager.RestartWorkloads(ctx, workloadNames, false)
 	if err != nil {
-		if errors.Is(err, workloads.ErrInvalidWorkloadName) {
+		if errors.Is(err, wt.ErrInvalidWorkloadName) {
 			http.Error(w, "Invalid workload name: "+err.Error(), http.StatusBadRequest)
 			return
 		}
@@ -457,7 +458,7 @@ func (s *WorkloadRoutes) deleteWorkloadsBulk(w http.ResponseWriter, r *http.Requ
 	// The request is not blocked on completion.
 	_, err = s.workloadManager.DeleteWorkloads(ctx, workloadNames)
 	if err != nil {
-		if errors.Is(err, workloads.ErrInvalidWorkloadName) {
+		if errors.Is(err, wt.ErrInvalidWorkloadName) {
 			http.Error(w, "Invalid workload name: "+err.Error(), http.StatusBadRequest)
 			return
 		}

pkg/runner/config.go

@@ -144,8 +144,7 @@ func (c *RunConfig) WriteJSON(w io.Writer) error {
 // ReadJSON deserializes the RunConfig from JSON read from the provided reader
 func ReadJSON(r io.Reader) (*RunConfig, error) {
 	var config RunConfig
-	decoder := json.NewDecoder(r)
-	if err := decoder.Decode(&config); err != nil {
+	if err := state.ReadJSON(r, &config); err != nil {
 		return nil, err
 	}
 	return &config, nil
@@ -313,38 +312,17 @@ func (c *RunConfig) WithStandardLabels() *RunConfig {
 	return c
 }
 
+// GetBaseName returns the base name for the run configuration
+func (c *RunConfig) GetBaseName() string {
+	return c.BaseName
+}
+
 // SaveState saves the run configuration to the state store
 func (c *RunConfig) SaveState(ctx context.Context) error {
-	// Create a state store
-	store, err := state.NewRunConfigStore(state.DefaultAppName)
-	if err != nil {
-		return fmt.Errorf("failed to create state store: %w", err)
-	}
-
-	// Get a writer for the state
-	writer, err := store.GetWriter(ctx, c.BaseName)
-	if err != nil {
-		return fmt.Errorf("failed to get writer for state: %w", err)
-	}
-	defer writer.Close()
-
-	// Serialize the configuration to JSON and write it directly to the state store
-	if err := c.WriteJSON(writer); err != nil {
-		return fmt.Errorf("failed to write run configuration: %w", err)
-	}
-
-	logger.Infof("Saved run configuration for %s", c.BaseName)
-	return nil
+	return state.SaveRunConfig(ctx, c)
 }
 
 // LoadState loads a run configuration from the state store
 func LoadState(ctx context.Context, name string) (*RunConfig, error) {
-	reader, err := state.LoadRunConfigJSON(ctx, name)
-	if err != nil {
-		return nil, err
-	}
-	defer reader.Close()
-
-	// Deserialize the configuration
-	return ReadJSON(reader)
+	return state.LoadRunConfig(ctx, name, ReadJSON)
 }

pkg/runner/runner.go

@@ -12,6 +12,7 @@ import (
 	"github.com/stacklok/toolhive/pkg/auth"
 	"github.com/stacklok/toolhive/pkg/client"
 	"github.com/stacklok/toolhive/pkg/config"
+	rt "github.com/stacklok/toolhive/pkg/container/runtime"
 	"github.com/stacklok/toolhive/pkg/labels"
 	"github.com/stacklok/toolhive/pkg/logger"
 	"github.com/stacklok/toolhive/pkg/mcp"
@@ -20,6 +21,7 @@ import (
 	"github.com/stacklok/toolhive/pkg/telemetry"
 	"github.com/stacklok/toolhive/pkg/transport"
 	"github.com/stacklok/toolhive/pkg/transport/types"
+	"github.com/stacklok/toolhive/pkg/workloads/statuses"
 )
 
 // Runner is responsible for running an MCP server with the provided configuration
@@ -32,12 +34,15 @@ type Runner struct {
 
 	// supportedMiddleware is a map of supported middleware types to their factory functions.
 	supportedMiddleware map[string]types.MiddlewareFactory
+
+	statusManager statuses.StatusManager
 }
 
 // NewRunner creates a new Runner with the provided configuration
-func NewRunner(runConfig *RunConfig) *Runner {
+func NewRunner(runConfig *RunConfig, statusManager statuses.StatusManager) *Runner {
 	return &Runner{
-		Config: runConfig,
+		Config:        runConfig,
+		statusManager: statusManager,
 	}
 }
 
@@ -294,6 +299,12 @@ func (r *Runner) Run(ctx context.Context) error {
 		}
 	}()
 
+	// At this point, we can consider the workload started successfully.
+	if err := r.statusManager.SetWorkloadStatus(ctx, r.Config.ContainerName, rt.WorkloadStatusRunning, ""); err != nil {
+		// If we can't set the status to `running` - treat it as a fatal error.
+		return fmt.Errorf("failed to set workload status: %v", err)
+	}
+
 	// Wait for either a signal or the done channel to be closed
 	select {
 	case sig := <-sigCh:

pkg/state/runconfig.go

@@ -2,6 +2,7 @@ package state
 
 import (
 	"context"
+	"encoding/json"
 	"fmt"
 	"io"
 
@@ -60,3 +61,89 @@ func DeleteSavedRunConfig(ctx context.Context, name string) error {
 	logger.Infof("Deleted run configuration for %s", name)
 	return nil
 }
+
+// RunConfigPersister defines an interface for objects that can be persisted and loaded as JSON
+type RunConfigPersister interface {
+	// WriteJSON serializes the object to JSON and writes it to the provided writer
+	WriteJSON(w io.Writer) error
+	// GetBaseName returns the base name used for persistence
+	GetBaseName() string
+}
+
+// ReadJSONFunc defines a function type for reading JSON into an object
+type ReadJSONFunc[T any] func(r io.Reader) (T, error)
+
+// SaveRunConfig saves a run configuration to the state store
+func SaveRunConfig[T RunConfigPersister](ctx context.Context, config T) error {
+	// Create a state store
+	store, err := NewRunConfigStore(DefaultAppName)
+	if err != nil {
+		return fmt.Errorf("failed to create state store: %w", err)
+	}
+
+	// Get a writer for the state
+	writer, err := store.GetWriter(ctx, config.GetBaseName())
+	if err != nil {
+		return fmt.Errorf("failed to get writer for state: %w", err)
+	}
+	defer writer.Close()
+
+	// Serialize the configuration to JSON and write it directly to the state store
+	if err := config.WriteJSON(writer); err != nil {
+		return fmt.Errorf("failed to write run configuration: %w", err)
+	}
+
+	logger.Infof("Saved run configuration for %s", config.GetBaseName())
+	return nil
+}
+
+// LoadRunConfig loads a run configuration from the state store using the provided reader function
+func LoadRunConfig[T any](ctx context.Context, name string, readJSONFunc ReadJSONFunc[T]) (T, error) {
+	var zero T
+	reader, err := LoadRunConfigJSON(ctx, name)
+	if err != nil {
+		return zero, err
+	}
+	defer reader.Close()
+
+	// Deserialize the configuration using the provided function
+	return readJSONFunc(reader)
+}
+
+// ReadRunConfigJSON deserializes a run configuration from JSON read from the provided reader
+// This is a generic JSON deserializer for any type that can be unmarshalled from JSON
+func ReadRunConfigJSON[T any](r io.Reader) (*T, error) {
+	var config T
+	decoder := json.NewDecoder(r)
+	if err := decoder.Decode(&config); err != nil {
+		return nil, err
+	}
+	return &config, nil
+}
+
+// LoadRunConfigOfType loads a run configuration of a specific type T from the state store
+func LoadRunConfigOfType[T any](ctx context.Context, name string) (*T, error) {
+	return LoadRunConfig(ctx, name, ReadRunConfigJSON[T])
+}
+
+// RunConfigReadJSONFunc defines the function signature for reading a RunConfig from JSON
+// This allows us to accept the runner.ReadJSON function without creating a circular dependency
+type RunConfigReadJSONFunc func(r io.Reader) (interface{}, error)
+
+// LoadRunConfigWithFunc loads a run configuration using a provided read function
+func LoadRunConfigWithFunc(ctx context.Context, name string, readFunc RunConfigReadJSONFunc) (interface{}, error) {
+	reader, err := LoadRunConfigJSON(ctx, name)
+	if err != nil {
+		return nil, err
+	}
+	defer reader.Close()
+
+	return readFunc(reader)
+}
+
+// ReadJSON deserializes JSON from the provided reader into a generic interface
+// This function is moved from the runner package to avoid circular dependencies
+func ReadJSON(r io.Reader, target interface{}) error {
+	decoder := json.NewDecoder(r)
+	return decoder.Decode(target)
+}