Fix env variables for MCP Gateway

Signed-off-by: David Gageot <david.gageot@docker.com>

Author: dgageot

pkg/config/gather.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"errors"
 	"fmt"
+	"os"
 	"sort"
 	"strings"
 
@@ -85,47 +86,41 @@ func GatherEnvVarsForTools(ctx context.Context, cfg *latest.Config) ([]string, e
 	requiredEnv := map[string]bool{}
 	var errs []error
 
-	for _, ref := range gatherMCPServerReferences(cfg) {
-		mcpServerName := gateway.ParseServerRef(ref)
-
-		secrets, err := gateway.RequiredEnvVars(ctx, mcpServerName)
-		if err != nil {
-			errs = append(errs, fmt.Errorf("reading which secrets the MCP server needs for %s: %w", ref, err))
-			continue
-		}
-
-		for _, secret := range secrets {
-			requiredEnv[secret.Env] = true
-		}
-	}
-
-	if len(errs) > 0 {
-		return mcpToSortedList(requiredEnv), fmt.Errorf("tool env preflight: %w", errors.Join(errs...))
-	}
-	return mcpToSortedList(requiredEnv), nil
-}
-
-func gatherMCPServerReferences(cfg *latest.Config) []string {
-	servers := map[string]bool{}
-
 	for i := range cfg.Agents {
 		agent := cfg.Agents[i]
+
 		for j := range agent.Toolsets {
 			toolSet := agent.Toolsets[j]
+			ref := toolSet.Ref
+			if toolSet.Type != "mcp" || ref == "" {
+				continue
+			}
 
-			if toolSet.Type == "mcp" && toolSet.Ref != "" {
-				servers[toolSet.Ref] = true
+			mcpServerName := gateway.ParseServerRef(ref)
+			secrets, err := gateway.RequiredEnvVars(ctx, mcpServerName)
+			if err != nil {
+				errs = append(errs, fmt.Errorf("reading which secrets the MCP server needs for %s: %w", ref, err))
+				continue
+			}
+
+			for _, secret := range secrets {
+				value, ok := toolSet.Env[secret.Env]
+				if !ok {
+					requiredEnv[secret.Env] = true
+				} else {
+					os.Expand(value, func(name string) string {
+						requiredEnv[name] = true
+						return ""
+					})
+				}
 			}
 		}
 	}
 
-	var list []string
-	for e := range servers {
-		list = append(list, e)
+	if len(errs) > 0 {
+		return mcpToSortedList(requiredEnv), fmt.Errorf("tool env preflight: %w", errors.Join(errs...))
 	}
-	sort.Strings(list)
-
-	return list
+	return mcpToSortedList(requiredEnv), nil
 }
 
 func mcpToSortedList(requiredEnv map[string]bool) []string {

pkg/environment/env.go

@@ -3,6 +3,7 @@ package environment
 import (
 	"context"
 	"os"
+	"strings"
 )
 
 // OsEnvProvider provides access to the operating system's environment variables.
@@ -16,22 +17,44 @@ func (p *OsEnvProvider) Get(_ context.Context, name string) string {
 	return os.Getenv(name)
 }
 
-type EnvFilesProviders struct {
+// EnvListProvider provides access a list of environment variables.
+type EnvListProvider struct {
+	env []string
+}
+
+func NewEnvListProvider(env []string) *EnvListProvider {
+	return &EnvListProvider{
+		env: env,
+	}
+}
+
+func (p *EnvListProvider) Get(_ context.Context, name string) string {
+	for _, e := range p.env {
+		n, v, ok := strings.Cut(e, "=")
+		if ok && n == name {
+			return v
+		}
+	}
+	return ""
+}
+
+// EnvFilesProvider provides access env files.
+type EnvFilesProvider struct {
 	values []KeyValuePair
 }
 
-func NewEnvFilesProvider(absEnvFiles []string) (*EnvFilesProviders, error) {
+func NewEnvFilesProvider(absEnvFiles []string) (*EnvFilesProvider, error) {
 	values, err := ReadEnvFiles(absEnvFiles)
 	if err != nil {
 		return nil, err
 	}
 
-	return &EnvFilesProviders{
+	return &EnvFilesProvider{
 		values: values,
 	}, nil
 }
 
-func (p *EnvFilesProviders) Get(_ context.Context, name string) string {
+func (p *EnvFilesProvider) Get(_ context.Context, name string) string {
 	for _, kv := range p.values {
 		if kv.Key == name {
 			return kv.Value

pkg/environment/env_test.go

@@ -21,3 +21,21 @@ func TestOsEnvProvider(t *testing.T) {
 	value = provider.Get(t.Context(), "NOT_FOUND")
 	assert.Empty(t, value)
 }
+
+func TestNewEnvListProvider(t *testing.T) {
+	t.Parallel()
+
+	provider := NewEnvListProvider([]string{
+		"TEST1=VALUE1",
+		"TEST2=VALUE2",
+	})
+
+	value := provider.Get(t.Context(), "TEST1")
+	assert.Equal(t, "VALUE1", value)
+
+	value = provider.Get(t.Context(), "TEST2")
+	assert.Equal(t, "VALUE2", value)
+
+	value = provider.Get(t.Context(), "NOT_FOUND")
+	assert.Empty(t, value)
+}

pkg/teamloader/registry.go

@@ -198,7 +198,17 @@ func createMCPTool(ctx context.Context, toolset latest.Toolset, _ string, runCon
 			return mcp.NewRemoteToolset(toolset.Name, serverSpec.Remote.URL, serverSpec.Remote.TransportType, nil), nil
 		}
 
-		return mcp.NewGatewayToolset(ctx, toolset.Name, mcpServerName, toolset.Config, runConfig.EnvProvider(), runConfig.WorkingDir)
+		env, err := environment.ExpandAll(ctx, environment.ToValues(toolset.Env), runConfig.EnvProvider())
+		if err != nil {
+			return nil, fmt.Errorf("failed to expand the tool's environment variables: %w", err)
+		}
+
+		envProvider := environment.NewMultiProvider(
+			environment.NewEnvListProvider(env),
+			runConfig.EnvProvider(),
+		)
+
+		return mcp.NewGatewayToolset(ctx, toolset.Name, mcpServerName, toolset.Config, envProvider, runConfig.WorkingDir)
 	}
 
 	if toolset.Command != "" {