Consistently use maps instead of slices for env vars (#1398)

This change is intended to allow env vars to be specified to the API as a map of string to string. Some refactoring was made to apply this change consistently throughout the codebase. Now, the CLI is responsible for converting from a list of strings to the map.

Also replace a manual mock with one generated by mockgen.

Author: dmjb

cmd/thv-proxyrunner/app/run.go

@@ -8,6 +8,7 @@ import (
 	"github.com/spf13/cobra"
 
 	"github.com/stacklok/toolhive/pkg/container"
+	"github.com/stacklok/toolhive/pkg/environment"
 	"github.com/stacklok/toolhive/pkg/logger"
 	"github.com/stacklok/toolhive/pkg/registry"
 	"github.com/stacklok/toolhive/pkg/runner"
@@ -237,6 +238,12 @@ func runCmdFunc(cmd *cobra.Command, args []string) error {
 
 	var imageMetadata *registry.ImageMetadata
 
+	// Parse environment variables from slice to map
+	envVarsMap, err := environment.ParseEnvironmentVariables(runEnv)
+	if err != nil {
+		return fmt.Errorf("failed to parse environment variables: %v", err)
+	}
+
 	// Initialize a new RunConfig with values from command-line flags
 	runConfig, err := runner.NewRunConfigBuilder().
 		WithRuntime(rt).
@@ -261,7 +268,7 @@ func runCmdFunc(cmd *cobra.Command, args []string) error {
 		WithTelemetryConfig(finalOtelEndpoint, runOtelEnablePrometheusMetricsPath, runOtelServiceName,
 			finalOtelSamplingRate, runOtelHeaders, runOtelInsecure, finalOtelEnvironmentVariables).
 		WithToolsFilter(runToolsFilter).
-		Build(ctx, imageMetadata, runEnv, envVarValidator)
+		Build(ctx, imageMetadata, envVarsMap, envVarValidator)
 	if err != nil {
 		return fmt.Errorf("failed to create RunConfig: %v", err)
 	}

cmd/thv/app/mcp_serve_helpers.go

@@ -66,11 +66,11 @@ func buildServerConfig(
 	builder = builder.WithTransportAndPorts(transport, 0, 0)
 
 	// Prepare environment variables
-	envVarsList := prepareEnvironmentVariables(imageMetadata, args.Env)
+	envVars := prepareEnvironmentVariables(imageMetadata, args.Env)
 
 	// Build the configuration
 	envVarValidator := &runner.DetachedEnvVarValidator{}
-	return builder.Build(ctx, imageMetadata, envVarsList, envVarValidator)
+	return builder.Build(ctx, imageMetadata, envVars, envVarValidator)
 }
 
 // configureTransport sets up transport configuration from metadata
@@ -88,7 +88,7 @@ func configureTransport(builder *runner.RunConfigBuilder, imageMetadata *registr
 }
 
 // prepareEnvironmentVariables merges default and user environment variables
-func prepareEnvironmentVariables(imageMetadata *registry.ImageMetadata, userEnv map[string]string) []string {
+func prepareEnvironmentVariables(imageMetadata *registry.ImageMetadata, userEnv map[string]string) map[string]string {
 	envVarsMap := make(map[string]string)
 
 	// Add default environment variables from metadata
@@ -105,13 +105,7 @@ func prepareEnvironmentVariables(imageMetadata *registry.ImageMetadata, userEnv
 		envVarsMap[k] = v
 	}
 
-	// Convert map to []string format
-	var envVarsList []string
-	for k, v := range envVarsMap {
-		envVarsList = append(envVarsList, fmt.Sprintf("%s=%s", k, v))
-	}
-
-	return envVarsList
+	return envVarsMap
 }
 
 // saveAndRunServer saves the configuration and runs the server

cmd/thv/app/run_flags.go

@@ -9,6 +9,7 @@ import (
 	cfg "github.com/stacklok/toolhive/pkg/config"
 	"github.com/stacklok/toolhive/pkg/container"
 	"github.com/stacklok/toolhive/pkg/container/runtime"
+	"github.com/stacklok/toolhive/pkg/environment"
 	"github.com/stacklok/toolhive/pkg/ignore"
 	"github.com/stacklok/toolhive/pkg/networking"
 	"github.com/stacklok/toolhive/pkg/process"
@@ -184,16 +185,16 @@ func AddRunFlags(cmd *cobra.Command, config *RunFlags) {
 // BuildRunnerConfig creates a runner.RunConfig from the configuration
 func BuildRunnerConfig(
 	ctx context.Context,
-	runConfig *RunFlags,
+	runFlags *RunFlags,
 	serverOrImage string,
 	cmdArgs []string,
 	debugMode bool,
 	cmd *cobra.Command,
 ) (*runner.RunConfig, error) {
 	// Validate the host flag
-	validatedHost, err := ValidateAndNormaliseHostFlag(runConfig.Host)
+	validatedHost, err := ValidateAndNormaliseHostFlag(runFlags.Host)
 	if err != nil {
-		return nil, fmt.Errorf("invalid host: %s", runConfig.Host)
+		return nil, fmt.Errorf("invalid host: %s", runFlags.Host)
 	}
 
 	// Get OIDC flags
@@ -212,15 +213,15 @@ func BuildRunnerConfig(
 	// Get OTEL flag values with config fallbacks
 	config := cfg.GetConfig()
 	finalOtelEndpoint, finalOtelSamplingRate, finalOtelEnvironmentVariables := getTelemetryFromFlags(cmd, config,
-		runConfig.OtelEndpoint, runConfig.OtelSamplingRate, runConfig.OtelEnvironmentVariables)
+		runFlags.OtelEndpoint, runFlags.OtelSamplingRate, runFlags.OtelEnvironmentVariables)
 
 	// Create container runtime
 	rt, err := container.NewFactory().Create(ctx)
 	if err != nil {
 		return nil, fmt.Errorf("failed to create container runtime: %v", err)
 	}
 
-	// Select an env var validation strategy depending on how the CLI is run:
+	// Select an envVars var validation strategy depending on how the CLI is run:
 	// If we have called the CLI directly, we use the CLIEnvVarValidator.
 	// If we are running in detached mode, or the CLI is wrapped by the K8s operator,
 	// we use the DetachedEnvVarValidator.
@@ -241,52 +242,58 @@ func BuildRunnerConfig(
 		// Take the MCP server we were supplied and either fetch the image, or
 		// build it from a protocol scheme. If the server URI refers to an image
 		// in our trusted registry, we will also fetch the image metadata.
-		imageURL, imageMetadata, err = retriever.GetMCPServer(ctx, serverOrImage, runConfig.CACertPath, runConfig.VerifyImage)
+		imageURL, imageMetadata, err = retriever.GetMCPServer(ctx, serverOrImage, runFlags.CACertPath, runFlags.VerifyImage)
 		if err != nil {
 			return nil, fmt.Errorf("failed to find or create the MCP server %s: %v", serverOrImage, err)
 		}
 	}
 
 	// Validate proxy mode early
-	if !types.IsValidProxyMode(runConfig.ProxyMode) {
-		if runConfig.ProxyMode == "" {
-			runConfig.ProxyMode = types.ProxyModeSSE.String() // default to SSE for backward compatibility
+	if !types.IsValidProxyMode(runFlags.ProxyMode) {
+		if runFlags.ProxyMode == "" {
+			runFlags.ProxyMode = types.ProxyModeSSE.String() // default to SSE for backward compatibility
 		} else {
-			return nil, fmt.Errorf("invalid value for --proxy-mode: %s", runConfig.ProxyMode)
+			return nil, fmt.Errorf("invalid value for --proxy-mode: %s", runFlags.ProxyMode)
 		}
 	}
 
+	// Parse the environment variables from a list of strings to a map.
+	envVars, err := environment.ParseEnvironmentVariables(runFlags.Env)
+	if err != nil {
+		return nil, fmt.Errorf("failed to parse environment variables: %v", err)
+	}
+
 	// Initialize a new RunConfig with values from command-line flags
 	return runner.NewRunConfigBuilder().
 		WithRuntime(rt).
 		WithCmdArgs(cmdArgs).
-		WithName(runConfig.Name).
+		WithName(runFlags.Name).
 		WithImage(imageURL).
 		WithHost(validatedHost).
-		WithTargetHost(runConfig.TargetHost).
+		WithTargetHost(runFlags.TargetHost).
 		WithDebug(debugMode).
-		WithVolumes(runConfig.Volumes).
-		WithSecrets(runConfig.Secrets).
-		WithAuthzConfigPath(runConfig.AuthzConfig).
-		WithAuditConfigPath(runConfig.AuditConfig).
-		WithPermissionProfileNameOrPath(runConfig.PermissionProfile).
-		WithNetworkIsolation(runConfig.IsolateNetwork).
-		WithK8sPodPatch(runConfig.K8sPodPatch).
-		WithProxyMode(types.ProxyMode(runConfig.ProxyMode)).
-		WithTransportAndPorts(runConfig.Transport, runConfig.ProxyPort, runConfig.TargetPort).
-		WithAuditEnabled(runConfig.EnableAudit, runConfig.AuditConfig).
-		WithLabels(runConfig.Labels).
-		WithGroup(runConfig.Group).
+		WithVolumes(runFlags.Volumes).
+		WithSecrets(runFlags.Secrets).
+		WithAuthzConfigPath(runFlags.AuthzConfig).
+		WithAuditConfigPath(runFlags.AuditConfig).
+		WithPermissionProfileNameOrPath(runFlags.PermissionProfile).
+		WithNetworkIsolation(runFlags.IsolateNetwork).
+		WithK8sPodPatch(runFlags.K8sPodPatch).
+		WithProxyMode(types.ProxyMode(runFlags.ProxyMode)).
+		WithTransportAndPorts(runFlags.Transport, runFlags.ProxyPort, runFlags.TargetPort).
+		WithAuditEnabled(runFlags.EnableAudit, runFlags.AuditConfig).
+		WithLabels(runFlags.Labels).
+		WithGroup(runFlags.Group).
 		WithOIDCConfig(oidcIssuer, oidcAudience, oidcJwksURL, oidcIntrospectionURL, oidcClientID, oidcClientSecret,
-			runConfig.ThvCABundle, runConfig.JWKSAuthTokenFile, runConfig.ResourceURL, runConfig.JWKSAllowPrivateIP).
-		WithTelemetryConfig(finalOtelEndpoint, runConfig.OtelEnablePrometheusMetricsPath, runConfig.OtelServiceName,
-			finalOtelSamplingRate, runConfig.OtelHeaders, runConfig.OtelInsecure, finalOtelEnvironmentVariables).
-		WithToolsFilter(runConfig.ToolsFilter).
+			runFlags.ThvCABundle, runFlags.JWKSAuthTokenFile, runFlags.ResourceURL, runFlags.JWKSAllowPrivateIP).
+		WithTelemetryConfig(finalOtelEndpoint, runFlags.OtelEnablePrometheusMetricsPath, runFlags.OtelServiceName,
+			finalOtelSamplingRate, runFlags.OtelHeaders, runFlags.OtelInsecure, finalOtelEnvironmentVariables).
+		WithToolsFilter(runFlags.ToolsFilter).
 		WithIgnoreConfig(&ignore.Config{
-			LoadGlobal:    runConfig.IgnoreGlobally,
-			PrintOverlays: runConfig.PrintOverlays,
+			LoadGlobal:    runFlags.IgnoreGlobally,
+			PrintOverlays: runFlags.PrintOverlays,
 		}).
-		Build(ctx, imageMetadata, runConfig.Env, envVarValidator)
+		Build(ctx, imageMetadata, envVars, envVarValidator)
 }
 
 // getOidcFromFlags extracts OIDC configuration from command flags

docs/server/docs.go

@@ -562,7 +562,7 @@ type createRequest struct {
 	// Port to expose from the container
 	TargetPort int `json:"target_port"`
 	// Environment variables to set in the container
-	EnvVars []string `json:"env_vars"`
+	EnvVars map[string]string `json:"env_vars"`
 	// Secret parameters to inject
 	Secrets []secrets.SecretParameter `json:"secrets"`
 	// Volume mounts

docs/server/swagger.json

@@ -242,19 +242,14 @@ func (c *RunConfig) WithPorts(proxyPort, targetPort int) (*RunConfig, error) {
 	return c, nil
 }
 
-// WithEnvironmentVariables parses and sets environment variables
-func (c *RunConfig) WithEnvironmentVariables(envVarStrings []string) (*RunConfig, error) {
-	envVars, err := environment.ParseEnvironmentVariables(envVarStrings)
-	if err != nil {
-		return c, fmt.Errorf("failed to parse environment variables: %v", err)
-	}
-
+// WithEnvironmentVariables sets environment variables
+func (c *RunConfig) WithEnvironmentVariables(envVars map[string]string) (*RunConfig, error) {
 	// Initialize EnvVars if it's nil
 	if c.EnvVars == nil {
 		c.EnvVars = make(map[string]string)
 	}
 
-	// Merge the parsed environment variables with existing ones
+	// Merge the provided environment variables with existing ones
 	for key, value := range envVars {
 		c.EnvVars[key] = value
 	}

docs/server/swagger.yaml

@@ -297,8 +297,12 @@ func (b *RunConfigBuilder) WithIgnoreConfig(ignoreConfig *ignore.Config) *RunCon
 }
 
 // Build creates the final RunConfig instance with validation and processing
-func (b *RunConfigBuilder) Build(ctx context.Context, imageMetadata *registry.ImageMetadata,
-	envVars []string, envVarValidator EnvVarValidator) (*RunConfig, error) {
+func (b *RunConfigBuilder) Build(
+	ctx context.Context,
+	imageMetadata *registry.ImageMetadata,
+	envVars map[string]string,
+	envVarValidator EnvVarValidator,
+) (*RunConfig, error) {
 	// When using the CLI validation strategy, this is where the prompting for
 	// missing environment variables will happen.
 	processedEnvVars, err := envVarValidator.Validate(ctx, imageMetadata, b.config, envVars)