feat: adds auto-remove flag for run (#115)

Signed-off-by: ChrisJBurns <29541485+ChrisJBurns@users.noreply.github.com>

Author: ChrisJBurns

cmd/thv/run.go

@@ -40,9 +40,11 @@ var (
 	runVolumes           []string
 	runSecrets           []string
 	runAuthzConfig       string
+	autoRemove           bool
 )
 
 func init() {
+	runCmd.Flags().BoolVar(&autoRemove, "rm", false, "Auto-remove the server container when the server has been stopped")
 	runCmd.Flags().StringVar(&runTransport, "transport", "stdio", "Transport mode (sse or stdio)")
 	runCmd.Flags().StringVar(&runName, "name", "", "Name of the MCP server (auto-generated from image if not provided)")
 	runCmd.Flags().IntVar(&runPort, "port", 0, "Port for the HTTP proxy to listen on (host port)")
@@ -120,6 +122,7 @@ func runCmdFunc(cmd *cobra.Command, args []string) error {
 		cmdArgs,
 		runName,
 		debugMode,
+		autoRemove,
 		runVolumes,
 		runSecrets,
 		runAuthzConfig,

cmd/thv/run_common.go

@@ -67,6 +67,14 @@ func detachProcess(cmd *cobra.Command, options *runner.RunConfig) error {
 		detachedArgs = append(detachedArgs, "--transport", string(options.Transport))
 	}
 
+	if options.Debug {
+		detachedArgs = append(detachedArgs, "--debug")
+	}
+
+	if options.AutoRemove {
+		detachedArgs = append(detachedArgs, "--rm")
+	}
+
 	// Use Name if available
 	if options.Name != "" {
 		detachedArgs = append(detachedArgs, "--name", options.Name)

pkg/container/docker/client.go

@@ -238,6 +238,7 @@ func (c *Client) CreateContainer(
 
 	// Create host configuration
 	hostConfig := &container.HostConfig{
+		AutoRemove:  false,
 		Mounts:      convertMounts(permissionConfig.Mounts),
 		NetworkMode: container.NetworkMode(permissionConfig.NetworkMode),
 		CapAdd:      permissionConfig.CapAdd,

pkg/runner/config.go

@@ -62,6 +62,9 @@ type RunConfig struct {
 	// Debug indicates whether debug mode is enabled
 	Debug bool `json:"debug,omitempty" yaml:"debug,omitempty"`
 
+	// AutoRemove indicates whether the container should be removed when the server has been stopped
+	AutoRemove bool `json:"auto_remove,omitempty" yaml:"auto_remove,omitempty"`
+
 	// Volumes are the directory mounts to pass to the container
 	// Format: "host-path:container-path[:ro]"
 	Volumes []string `json:"volumes,omitempty" yaml:"volumes,omitempty"`
@@ -117,6 +120,7 @@ func NewRunConfigFromFlags(
 	cmdArgs []string,
 	name string,
 	debug bool,
+	autoRemove bool,
 	volumes []string,
 	secretsList []string,
 	authzConfigPath string,
@@ -132,6 +136,7 @@ func NewRunConfigFromFlags(
 		CmdArgs:                     cmdArgs,
 		Name:                        name,
 		Debug:                       debug,
+		AutoRemove:                  autoRemove,
 		Volumes:                     volumes,
 		Secrets:                     secretsList,
 		AuthzConfigPath:             authzConfigPath,

pkg/runner/config_test.go

@@ -801,6 +801,7 @@ func TestNewRunConfigFromFlags(t *testing.T) {
 	cmdArgs := []string{"arg1", "arg2"}
 	name := "test-server"
 	debug := true
+	autoRemove := true
 	volumes := []string{"/host:/container"}
 	secretsList := []string{"secret1,target=ENV_VAR1"}
 	authzConfigPath := "/path/to/authz.json"
@@ -816,6 +817,7 @@ func TestNewRunConfigFromFlags(t *testing.T) {
 		cmdArgs,
 		name,
 		debug,
+		autoRemove,
 		volumes,
 		secretsList,
 		authzConfigPath,

pkg/runner/runner.go

@@ -43,6 +43,7 @@ func (r *Runner) Run(ctx context.Context) error {
 		TargetHost: r.Config.TargetHost,
 		Runtime:    r.Config.Runtime,
 		Debug:      r.Config.Debug,
+		AutoRemove: r.Config.AutoRemove,
 	}
 
 	// Add OIDC middleware if OIDC validation is enabled

pkg/transport/factory.go

@@ -19,14 +19,15 @@ func NewFactory() *Factory {
 func (*Factory) Create(config types.Config) (types.Transport, error) {
 	switch config.Type {
 	case types.TransportTypeStdio:
-		return NewStdioTransport(config.Port, config.Runtime, config.Debug, config.Middlewares...), nil
+		return NewStdioTransport(config.Port, config.Runtime, config.Debug, config.AutoRemove, config.Middlewares...), nil
 	case types.TransportTypeSSE:
 		return NewSSETransport(
 			config.Host,
 			config.Port,
 			config.TargetPort,
 			config.Runtime,
 			config.Debug,
+			config.AutoRemove,
 			config.TargetHost,
 			config.Middlewares...,
 		), nil

pkg/transport/sse.go

@@ -30,6 +30,7 @@ type SSETransport struct {
 	containerName string
 	runtime       rt.Runtime
 	debug         bool
+	autoRemove    bool
 	middlewares   []types.Middleware
 
 	// Mutex for protecting shared state
@@ -53,6 +54,7 @@ func NewSSETransport(
 	targetPort int,
 	runtime rt.Runtime,
 	debug bool,
+	autoRemove bool,
 	targetHost string,
 	middlewares ...types.Middleware,
 ) *SSETransport {
@@ -73,6 +75,7 @@ func NewSSETransport(
 		targetHost:  targetHost,
 		runtime:     runtime,
 		debug:       debug,
+		autoRemove:  autoRemove,
 		shutdownCh:  make(chan struct{}),
 	}
 }
@@ -253,15 +256,15 @@ func (t *SSETransport) Stop(ctx context.Context) error {
 			return fmt.Errorf("failed to stop container: %w", err)
 		}
 
-		// Remove the container if debug mode is not enabled
-		if !t.debug {
+		// Remove the container if auto-remove is enabled
+		if t.autoRemove {
 			logger.Log.Info(fmt.Sprintf("Removing container %s...", t.containerName))
 			if err := t.runtime.RemoveContainer(ctx, t.containerID); err != nil {
 				logger.Log.Warn(fmt.Sprintf("Warning: Failed to remove container: %v", err))
 			}
 			logger.Log.Info(fmt.Sprintf("Container %s removed", t.containerName))
 		} else {
-			logger.Log.Info(fmt.Sprintf("Debug mode enabled, container %s not removed", t.containerName))
+			logger.Log.Info(fmt.Sprintf("Auto-remove disabled, container %s not removed", t.containerName))
 		}
 	}
 

pkg/transport/stdio.go

@@ -29,6 +29,7 @@ type StdioTransport struct {
 	containerName string
 	runtime       rt.Runtime
 	debug         bool
+	autoRemove    bool
 	middlewares   []types.Middleware
 
 	// Mutex for protecting shared state
@@ -54,12 +55,14 @@ func NewStdioTransport(
 	port int,
 	runtime rt.Runtime,
 	debug bool,
+	autoRemove bool,
 	middlewares ...types.Middleware,
 ) *StdioTransport {
 	return &StdioTransport{
 		port:        port,
 		runtime:     runtime,
 		debug:       debug,
+		autoRemove:  autoRemove,
 		middlewares: middlewares,
 		shutdownCh:  make(chan struct{}),
 	}
@@ -242,16 +245,16 @@ func (t *StdioTransport) Stop(ctx context.Context) error {
 			}
 		}
 
-		// Remove the container if debug mode is not enabled
-		if !t.debug {
+		// Remove the container if auto-remove is enabled
+		if t.autoRemove {
 			logger.Log.Info(fmt.Sprintf("Removing container %s...", t.containerName))
 			if err := t.runtime.RemoveContainer(ctx, t.containerID); err != nil {
 				logger.Log.Error(fmt.Sprintf("Warning: Failed to remove container: %v", err))
 			} else {
 				logger.Log.Info(fmt.Sprintf("Container %s removed", t.containerName))
 			}
 		} else {
-			logger.Log.Info(fmt.Sprintf("Debug mode enabled, container %s not removed", t.containerName))
+			logger.Log.Info(fmt.Sprintf("Auto-remove disabled, container %s not removed", t.containerName))
 		}
 	}
 

pkg/transport/types/transport.go

@@ -123,6 +123,9 @@ type Config struct {
 	// If debug mode is enabled, containers will not be removed when stopped.
 	Debug bool
 
+	// AutoRemove indicates whether the container should be removed when the server has been stopped.
+	AutoRemove bool
+
 	// Middlewares is a list of middleware functions to apply to the transport.
 	// These are applied in order, with the first middleware being the outermost wrapper.
 	Middlewares []Middleware