Skip to content

Commit 541dce1

Browse files
dgageotdgageot
committed
Put back env vars at tools level
Signed-off-by: David Gageot <david.gageot@docker.com>
1 parent fd76c28 commit 541dce1

5 files changed

Lines changed: 85 additions & 39 deletions

File tree

‎pkg/config/v2/types.go‎

Lines changed: 8 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -62,13 +62,14 @@ type ScriptShellToolConfig struct {
6262

6363
// Toolset represents a tool configuration
6464
type Toolset struct {
65-
Type string `json:"type,omitempty" yaml:"type,omitempty"`
66-
Ref string `json:"ref,omitempty" yaml:"ref,omitempty"`
67-
Config any `json:"config,omitempty" yaml:"config,omitempty"`
68-
Command string `json:"command,omitempty" yaml:"command,omitempty"`
69-
Remote Remote `json:"remote,omitempty" yaml:"remote,omitempty"`
70-
Args []string `json:"args,omitempty" yaml:"args,omitempty"`
71-
Tools []string `json:"tools,omitempty" yaml:"tools,omitempty"`
65+
Type string `json:"type,omitempty" yaml:"type,omitempty"`
66+
Ref string `json:"ref,omitempty" yaml:"ref,omitempty"`
67+
Config any `json:"config,omitempty" yaml:"config,omitempty"`
68+
Command string `json:"command,omitempty" yaml:"command,omitempty"`
69+
Remote Remote `json:"remote,omitempty" yaml:"remote,omitempty"`
70+
Args []string `json:"args,omitempty" yaml:"args,omitempty"`
71+
Tools []string `json:"tools,omitempty" yaml:"tools,omitempty"`
72+
Env map[string]string `json:"env,omitempty" yaml:"env,omitempty"`
7273

7374
// For the think tool
7475
Shared bool `json:"shared,omitempty" yaml:"shared,omitempty"`

‎pkg/config/v2/upgrade.go‎

Lines changed: 0 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -24,9 +24,6 @@ func UpgradeFrom(old v1.Config) (Config, error) {
2424
for i := range agent.Toolsets {
2525
toolSet := agent.Toolsets[i]
2626

27-
if len(toolSet.Env) > 0 {
28-
return Config{}, errors.New("toolset Env is not supported anymore")
29-
}
3027
if len(toolSet.Envfiles) > 0 {
3128
return Config{}, errors.New("toolset Envfiles is not supported anymore")
3229
}

‎pkg/environment/env_files_test.go‎

Lines changed: 14 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -9,24 +9,23 @@ import (
99
"github.com/stretchr/testify/require"
1010
)
1111

12-
func TestExpand(t *testing.T) {
13-
env := []string{"USER=alice", "HOME=/home/alice"}
12+
func TestExpandAll(t *testing.T) {
13+
t.Setenv("USER", "alice")
14+
t.Setenv("HOME", "/home/alice")
1415

15-
tests := []struct {
16-
input string
17-
expected string
18-
}{
19-
{"Hello $USER", "Hello alice"},
20-
{"Your home is at $HOME", "Your home is at /home/alice"},
21-
{"No variable here", "No variable here"},
22-
{"$UNKNOWN_VAR should be empty", " should be empty"},
23-
}
16+
expanded, err := ExpandAll(t.Context(), []string{"Hello $USER", "Your home is at $HOME", "No variable here"}, NewOsEnvProvider())
2417

25-
for _, test := range tests {
26-
result := Expand(test.input, env)
18+
require.NoError(t, err)
19+
assert.Equal(t, []string{"Hello alice", "Your home is at /home/alice", "No variable here"}, expanded)
20+
}
2721

28-
assert.Equal(t, test.expected, result)
29-
}
22+
func TestExpandAll_Error(t *testing.T) {
23+
t.Setenv("UNKNOWN_VAR", "")
24+
25+
expanded, err := ExpandAll(t.Context(), []string{"$UNKNOWN_VAR"}, NewOsEnvProvider())
26+
27+
require.Error(t, err)
28+
assert.Empty(t, expanded)
3029
}
3130

3231
func TestAbsolutePath(t *testing.T) {

‎pkg/environment/expand.go‎

Lines changed: 40 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -1,17 +1,49 @@
11
package environment
22

33
import (
4+
"context"
5+
"fmt"
46
"os"
5-
"strings"
7+
"slices"
68
)
79

8-
func Expand(value string, env []string) string {
9-
return os.Expand(value, func(name string) string {
10-
for _, e := range env {
11-
if after, ok := strings.CutPrefix(e, name+"="); ok {
12-
return after
13-
}
10+
func ExpandAll(ctx context.Context, values []string, env Provider) ([]string, error) {
11+
var expandedEnv []string
12+
13+
for _, value := range values {
14+
expanded, err := Expand(ctx, value, env)
15+
if err != nil {
16+
return nil, err
17+
}
18+
19+
expandedEnv = append(expandedEnv, expanded)
20+
}
21+
22+
return expandedEnv, nil
23+
}
24+
25+
func Expand(ctx context.Context, value string, env Provider) (string, error) {
26+
var err error
27+
28+
expanded := os.Expand(value, func(name string) string {
29+
v := env.Get(ctx, name)
30+
if v == "" {
31+
err = fmt.Errorf("environment variable %q not set", name)
1432
}
15-
return ""
33+
return v
1634
})
35+
if err != nil {
36+
return "", err
37+
}
38+
39+
return expanded, nil
40+
}
41+
42+
func ToValues(envMap map[string]string) []string {
43+
var values []string
44+
for k, v := range envMap {
45+
values = append(values, fmt.Sprintf("%s=%s", k, v))
46+
}
47+
slices.Sort(values)
48+
return values
1749
}

‎pkg/teamloader/teamloader.go‎

Lines changed: 23 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -130,7 +130,7 @@ func Load(ctx context.Context, path string, runConfig config.RuntimeConfig) (*te
130130
if !ok {
131131
return nil, fmt.Errorf("agent '%s' not found in configuration", name)
132132
}
133-
agentTools, err := getToolsForAgent(&a, parentDir, sharedTools, models[0])
133+
agentTools, err := getToolsForAgent(ctx, &a, parentDir, sharedTools, models[0], env)
134134
if err != nil {
135135
return nil, fmt.Errorf("failed to get tools: %w", err)
136136
}
@@ -184,7 +184,7 @@ func getModelsForAgent(ctx context.Context, cfg *latest.Config, a *latest.AgentC
184184
}
185185

186186
// getToolsForAgent returns the tool definitions for an agent based on its configuration
187-
func getToolsForAgent(a *latest.AgentConfig, parentDir string, sharedTools map[string]tools.ToolSet, model provider.Provider) ([]tools.ToolSet, error) {
187+
func getToolsForAgent(ctx context.Context, a *latest.AgentConfig, parentDir string, sharedTools map[string]tools.ToolSet, model provider.Provider, env environment.Provider) ([]tools.ToolSet, error) {
188188
var t []tools.ToolSet
189189

190190
if len(a.SubAgents) > 0 {
@@ -278,17 +278,34 @@ func getToolsForAgent(a *latest.AgentConfig, parentDir string, sharedTools map[s
278278
// This improves shareability of agent configs.
279279
args = append(args, "--catalog=https://desktop.docker.com/mcp/catalog/v2/catalog.yaml")
280280

281+
envVars, err := environment.ExpandAll(ctx, environment.ToValues(toolset.Env), env)
282+
if err != nil {
283+
return nil, fmt.Errorf("failed to expand the tool's environment variables: %w", err)
284+
}
285+
281286
// TODO(dga): If the server's docker image had the right annotations, we could run it directly with `docker run` or with the MCP gateway as a go library.
282-
t = append(t, mcp.NewToolsetCommand("docker", args, os.Environ(), toolset.Tools, cleanUp))
287+
t = append(t, mcp.NewToolsetCommand("docker", args, envVars, toolset.Tools, cleanUp))
283288

284289
case toolset.Type == "mcp" && toolset.Command != "":
285-
t = append(t, mcp.NewToolsetCommand(toolset.Command, toolset.Args, os.Environ(), toolset.Tools, nil))
290+
envVars, err := environment.ExpandAll(ctx, environment.ToValues(toolset.Env), env)
291+
if err != nil {
292+
return nil, fmt.Errorf("failed to expand the tool's environment variables: %w", err)
293+
}
294+
295+
t = append(t, mcp.NewToolsetCommand(toolset.Command, toolset.Args, envVars, toolset.Tools, nil))
286296

287297
case toolset.Type == "mcp" && toolset.Remote.URL != "":
288-
// Expand headers.
298+
// TODO: the tool can have env variables too
299+
300+
// Expand env vars in headers.
289301
headers := map[string]string{}
290302
for k, v := range toolset.Remote.Headers {
291-
headers[k] = environment.Expand(v, os.Environ())
303+
expanded, err := environment.Expand(ctx, v, env)
304+
if err != nil {
305+
return nil, fmt.Errorf("failed to expand header '%s': %w", k, err)
306+
}
307+
308+
headers[k] = expanded
292309
}
293310

294311
mcpc, err := mcp.NewToolsetRemote(toolset.Remote.URL, toolset.Remote.TransportType, headers, toolset.Tools)

0 commit comments

Comments
 (0)