Merge pull request #584 from rumpl/move-things

Move things around

Author: rumpl

cmd/root/run.go

@@ -192,7 +192,7 @@ func doRunCommand(ctx context.Context, args []string, exec bool) error {
 			runConfig.RedirectURI = "http://localhost:8083/oauth-callback"
 		}
 
-		agents, err = teamloader.LoadWithOverrides(ctx, agentFilename, runConfig, modelOverrides)
+		agents, err = teamloader.Load(ctx, agentFilename, runConfig, teamloader.WithModelOverrides(modelOverrides))
 		if err != nil {
 			return err
 		}

pkg/config/config.go

@@ -1,6 +1,7 @@
 package config
 
 import (
+	"context"
 	"fmt"
 	"os"
 	"path/filepath"
@@ -11,6 +12,7 @@ import (
 	v0 "github.com/docker/cagent/pkg/config/v0"
 	v1 "github.com/docker/cagent/pkg/config/v1"
 	latest "github.com/docker/cagent/pkg/config/v2"
+	"github.com/docker/cagent/pkg/environment"
 	"github.com/docker/cagent/pkg/filesystem"
 )
 
@@ -55,6 +57,24 @@ func LoadConfig(path string, fs filesystem.FS) (*latest.Config, error) {
 	return &config, nil
 }
 
+// CheckRequiredEnvVars checks which environment variables are required by the models and tools.
+//
+// This allows exiting early with a proper error message instead of failing later when trying to use a model or tool.
+func CheckRequiredEnvVars(ctx context.Context, cfg *latest.Config, env environment.Provider, runtimeConfig RuntimeConfig) error {
+	requiredEnv, err := GatherMissingEnvVars(ctx, cfg, env, runtimeConfig)
+	if err != nil {
+		return fmt.Errorf("gathering required environment variables: %w", err)
+	}
+
+	if len(requiredEnv) == 0 {
+		return nil
+	}
+
+	return &environment.RequiredEnvError{
+		Missing: requiredEnv,
+	}
+}
+
 func parseCurrentVersion(dir string, data []byte, version any) (any, error) {
 	options := []yaml.DecodeOption{yaml.Strict(), yaml.ReferenceDirs(dir)}
 
@@ -142,58 +162,3 @@ func validateConfig(cfg *latest.Config) error {
 func boolPtr(b bool) *bool {
 	return &b
 }
-
-func ValidatePathInDirectory(path, allowedDir string) (string, error) {
-	if path == "" {
-		return "", fmt.Errorf("empty path")
-	}
-
-	cleanPath := filepath.Clean(path)
-
-	if cleanPath == "" || cleanPath == "." {
-		return "", fmt.Errorf("empty or invalid path")
-	}
-
-	if filepath.IsAbs(cleanPath) && allowedDir == "" {
-		if strings.Contains(path, "..") {
-			return "", fmt.Errorf("path contains directory traversal sequences")
-		}
-		return cleanPath, nil
-	}
-
-	if allowedDir == "" {
-		if strings.HasPrefix(cleanPath, "..") {
-			return "", fmt.Errorf("path contains directory traversal sequences")
-		}
-		return cleanPath, nil
-	}
-
-	cleanAllowedDir := filepath.Clean(allowedDir)
-	absAllowedDir, err := filepath.Abs(cleanAllowedDir)
-	if err != nil {
-		return "", fmt.Errorf("invalid allowed directory: %w", err)
-	}
-
-	var targetPath string
-	if filepath.IsAbs(cleanPath) {
-		targetPath = cleanPath
-	} else {
-		targetPath = filepath.Join(absAllowedDir, cleanPath)
-	}
-
-	absTargetPath, err := filepath.Abs(targetPath)
-	if err != nil {
-		return "", fmt.Errorf("invalid path: %w", err)
-	}
-
-	relPath, err := filepath.Rel(absAllowedDir, absTargetPath)
-	if err != nil {
-		return "", fmt.Errorf("cannot determine relative path: %w", err)
-	}
-
-	if strings.HasPrefix(relPath, "..") {
-		return "", fmt.Errorf("path outside allowed directory: %s", path)
-	}
-
-	return absTargetPath, nil
-}

pkg/config/config_test.go

@@ -1,11 +1,15 @@
 package config
 
 import (
+	"context"
 	"os"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
+
+	latest "github.com/docker/cagent/pkg/config/v2"
+	"github.com/docker/cagent/pkg/environment"
 )
 
 func TestAutoRegisterModels(t *testing.T) {
@@ -188,3 +192,235 @@ func openRoot(t *testing.T, dir string) *os.Root {
 
 	return root
 }
+
+type noEnvProvider struct{}
+
+func (p *noEnvProvider) Get(context.Context, string) string { return "" }
+
+func TestCheckRequiredEnvVars(t *testing.T) {
+	tests := []struct {
+		yaml            string
+		expectedMissing []string
+	}{
+		{
+			yaml:            "openai_inline.yaml",
+			expectedMissing: []string{"OPENAI_API_KEY"},
+		},
+		{
+			yaml:            "anthropic_inline.yaml",
+			expectedMissing: []string{"ANTHROPIC_API_KEY"},
+		},
+		{
+			yaml:            "google_inline.yaml",
+			expectedMissing: []string{"GOOGLE_API_KEY"},
+		},
+		{
+			yaml:            "dmr_inline.yaml",
+			expectedMissing: []string{},
+		},
+		{
+			yaml:            "openai_model.yaml",
+			expectedMissing: []string{"OPENAI_API_KEY"},
+		},
+		{
+			yaml:            "anthropic_model.yaml",
+			expectedMissing: []string{"ANTHROPIC_API_KEY"},
+		},
+		{
+			yaml:            "google_model.yaml",
+			expectedMissing: []string{"GOOGLE_API_KEY"},
+		},
+		{
+			yaml:            "dmr_model.yaml",
+			expectedMissing: []string{},
+		},
+		{
+			yaml:            "all.yaml",
+			expectedMissing: []string{"ANTHROPIC_API_KEY", "GOOGLE_API_KEY", "OPENAI_API_KEY"},
+		},
+	}
+	for _, test := range tests {
+		t.Run(test.yaml, func(t *testing.T) {
+			t.Parallel()
+
+			root := openRoot(t, "testdata/env")
+
+			cfg, err := LoadConfig(test.yaml, root)
+			require.NoError(t, err)
+
+			err = CheckRequiredEnvVars(t.Context(), cfg, &noEnvProvider{}, RuntimeConfig{})
+
+			if len(test.expectedMissing) == 0 {
+				require.NoError(t, err)
+			} else {
+				require.Error(t, err)
+				assert.Equal(t, test.expectedMissing, err.(*environment.RequiredEnvError).Missing)
+			}
+		})
+	}
+}
+
+func TestCheckRequiredEnvVarsWithModelGateway(t *testing.T) {
+	t.Parallel()
+
+	root := openRoot(t, "testdata/env")
+
+	cfg, err := LoadConfig("all.yaml", root)
+	require.NoError(t, err)
+
+	err = CheckRequiredEnvVars(t.Context(), cfg, &noEnvProvider{}, RuntimeConfig{
+		ModelsGateway: "gateway:8080",
+	})
+
+	require.NoError(t, err)
+}
+
+func TestApplyModelOverrides(t *testing.T) {
+	tests := []struct {
+		name        string
+		agents      map[string]latest.AgentConfig
+		overrides   []string
+		expected    map[string]string // agent name -> expected model
+		expectError bool
+		errorMsg    string
+	}{
+		{
+			name: "global override",
+			agents: map[string]latest.AgentConfig{
+				"root":  {Model: "openai/gpt-4"},
+				"other": {Model: "anthropic/claude-3"},
+			},
+			overrides: []string{"google/gemini-pro"},
+			expected: map[string]string{
+				"root":  "google/gemini-pro",
+				"other": "google/gemini-pro",
+			},
+		},
+		{
+			name: "single per-agent override",
+			agents: map[string]latest.AgentConfig{
+				"root":  {Model: "openai/gpt-4"},
+				"other": {Model: "anthropic/claude-3"},
+			},
+			overrides: []string{"other=google/gemini-pro"},
+			expected: map[string]string{
+				"root":  "openai/gpt-4",
+				"other": "google/gemini-pro",
+			},
+		},
+		{
+			name: "multiple separate flags",
+			agents: map[string]latest.AgentConfig{
+				"root":  {Model: "openai/gpt-4"},
+				"other": {Model: "anthropic/claude-3"},
+			},
+			overrides: []string{"root=openai/gpt-5", "other=anthropic/claude-sonnet-4-0"},
+			expected: map[string]string{
+				"root":  "openai/gpt-5",
+				"other": "anthropic/claude-sonnet-4-0",
+			},
+		},
+		{
+			name: "comma-separated format",
+			agents: map[string]latest.AgentConfig{
+				"root":  {Model: "openai/gpt-4"},
+				"other": {Model: "anthropic/claude-3"},
+				"third": {Model: "google/gemini-pro"},
+			},
+			overrides: []string{"root=openai/gpt-5,other=anthropic/claude-sonnet-4-0"},
+			expected: map[string]string{
+				"root":  "openai/gpt-5",
+				"other": "anthropic/claude-sonnet-4-0",
+				"third": "google/gemini-pro",
+			},
+		},
+		{
+			name: "mixed formats",
+			agents: map[string]latest.AgentConfig{
+				"root":     {Model: "openai/gpt-4"},
+				"other":    {Model: "anthropic/claude-3"},
+				"third":    {Model: "google/gemini-pro"},
+				"reviewer": {Model: "openai/gpt-3.5-turbo"},
+			},
+			overrides: []string{"root=openai/gpt-5,other=anthropic/claude-4", "reviewer=google/gemini-1.5-pro"},
+			expected: map[string]string{
+				"root":     "openai/gpt-5",
+				"other":    "anthropic/claude-4",
+				"third":    "google/gemini-pro",
+				"reviewer": "google/gemini-1.5-pro",
+			},
+		},
+		{
+			name: "last override wins",
+			agents: map[string]latest.AgentConfig{
+				"root": {Model: "openai/gpt-4"},
+			},
+			overrides: []string{"root=openai/gpt-5", "root=anthropic/claude-4"},
+			expected: map[string]string{
+				"root": "anthropic/claude-4",
+			},
+		},
+		{
+			name: "unknown agent error",
+			agents: map[string]latest.AgentConfig{
+				"root": {Model: "openai/gpt-4"},
+			},
+			overrides:   []string{"nonexistent=openai/gpt-5"},
+			expectError: true,
+			errorMsg:    "unknown agent 'nonexistent'",
+		},
+		{
+			name: "empty model spec error",
+			agents: map[string]latest.AgentConfig{
+				"root": {Model: "openai/gpt-4"},
+			},
+			overrides:   []string{"root="},
+			expectError: true,
+			errorMsg:    "empty model specification in override: root=",
+		},
+		{
+			name: "empty global model spec is skipped",
+			agents: map[string]latest.AgentConfig{
+				"root": {Model: "openai/gpt-4"},
+			},
+			overrides: []string{""},
+			expected: map[string]string{
+				"root": "openai/gpt-4",
+			},
+		},
+		{
+			name: "whitespace handling",
+			agents: map[string]latest.AgentConfig{
+				"root":  {Model: "openai/gpt-4"},
+				"other": {Model: "anthropic/claude-3"},
+			},
+			overrides: []string{" root = openai/gpt-5 , other = anthropic/claude-4 "},
+			expected: map[string]string{
+				"root":  "openai/gpt-5",
+				"other": "anthropic/claude-4",
+			},
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			t.Parallel()
+
+			cfg := &latest.Config{
+				Agents: tt.agents,
+				Models: make(map[string]latest.ModelConfig),
+			}
+
+			err := ApplyModelOverrides(cfg, tt.overrides)
+
+			if tt.expectError {
+				require.ErrorContains(t, err, tt.errorMsg)
+			} else {
+				require.NoError(t, err)
+				for agentName, expectedModel := range tt.expected {
+					assert.Equal(t, expectedModel, cfg.Agents[agentName].Model)
+				}
+			}
+		})
+	}
+}

pkg/secrets/gather.go pkg/config/gather.gopkg/secrets/gather.go renamed to pkg/config/gather.go

@@ -1,11 +1,10 @@
-package secrets
+package config
 
 import (
 	"context"
 	"fmt"
 	"sort"
 
-	"github.com/docker/cagent/pkg/config"
 	latest "github.com/docker/cagent/pkg/config/v2"
 	"github.com/docker/cagent/pkg/environment"
 	"github.com/docker/cagent/pkg/gateway"
@@ -15,7 +14,7 @@ import (
 // GatherMissingEnvVars finds out which environment variables are required by the models and tools.
 // This allows exiting early with a proper error message instead of failing later when trying to use a model or tool.
 // TODO(dga): This code contains lots of duplication and ought to be refactored.
-func GatherMissingEnvVars(ctx context.Context, cfg *latest.Config, env environment.Provider, runtimeConfig config.RuntimeConfig) ([]string, error) {
+func GatherMissingEnvVars(ctx context.Context, cfg *latest.Config, env environment.Provider, runtimeConfig RuntimeConfig) ([]string, error) {
 	requiredEnv := map[string]bool{}
 
 	// Models