Merge pull request #776 from dgageot/prepare-default-agent

Prepare for default agent

Author: dgageot

cmd/root/run.go

@@ -86,18 +86,21 @@ func (f *runExecFlags) runOrExec(ctx context.Context, out *cli.Printer, args []s
 	var rt runtime.Runtime
 	var sess *session.Session
 	var err error
-	if f.remoteAddress != "" {
-		rt, sess, err = f.createRemoteRuntimeAndSession(ctx, args[0])
+	switch {
+	case f.remoteAddress != "":
+		agentFileName = args[0]
+		rt, sess, err = f.createRemoteRuntimeAndSession(ctx, agentFileName)
 		if err != nil {
 			return err
 		}
-	} else {
+
+	default:
 		agentFileName, err = f.resolveAgentFile(ctx, out, args[0])
 		if err != nil {
 			return err
 		}
 
-		t, err := f.loadAgents(ctx, agentFileName)
+		t, err := f.loadAgentFrom(ctx, teamloader.NewFileSource(agentFileName))
 		if err != nil {
 			return err
 		}
@@ -133,8 +136,8 @@ func (f *runExecFlags) resolveAgentFile(ctx context.Context, out *cli.Printer, a
 	return agentfile.Resolve(ctx, out, agentFilename)
 }
 
-func (f *runExecFlags) loadAgents(ctx context.Context, agentFilename string) (*team.Team, error) {
-	t, err := teamloader.Load(ctx, agentFilename, f.runConfig, teamloader.WithModelOverrides(f.modelOverrides))
+func (f *runExecFlags) loadAgentFrom(ctx context.Context, source teamloader.AgentSource) (*team.Team, error) {
+	t, err := teamloader.LoadFrom(ctx, source, f.runConfig, teamloader.WithModelOverrides(f.modelOverrides))
 	if err != nil {
 		return nil, err
 	}

pkg/config/config.go

@@ -19,22 +19,26 @@ import (
 func LoadConfig(path string, fs filesystem.FS) (*latest.Config, error) {
 	data, err := fs.ReadFile(path)
 	if err != nil {
-		return nil, fmt.Errorf("reading config file: %w", err)
+		return nil, fmt.Errorf("reading config file %s: %w", path, err)
 	}
 
+	return LoadConfigBytes(data)
+}
+
+func LoadConfigBytes(data []byte) (*latest.Config, error) {
 	var raw struct {
 		Version string `yaml:"version,omitempty"`
 	}
 	if err := yaml.UnmarshalWithOptions(data, &raw); err != nil {
-		return nil, fmt.Errorf("looking for version in config file %s\n%s", path, yaml.FormatError(err, true, true))
+		return nil, fmt.Errorf("looking for version in config file\n%s", yaml.FormatError(err, true, true))
 	}
 	if raw.Version == "" {
 		raw.Version = latest.Version
 	}
 
 	oldConfig, err := parseCurrentVersion(data, raw.Version)
 	if err != nil {
-		return nil, fmt.Errorf("parsing config file %s\n%s", path, yaml.FormatError(err, true, true))
+		return nil, fmt.Errorf("parsing config file\n%s", yaml.FormatError(err, true, true))
 	}
 
 	config, err := migrateToLatestConfig(oldConfig)

pkg/teamloader/registry.go

@@ -0,0 +1,225 @@
+package teamloader
+
+import (
+	"context"
+	"fmt"
+	"os"
+	"path/filepath"
+	"time"
+
+	"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"
+	"github.com/docker/cagent/pkg/js"
+	"github.com/docker/cagent/pkg/memory/database/sqlite"
+	"github.com/docker/cagent/pkg/path"
+	"github.com/docker/cagent/pkg/tools"
+	"github.com/docker/cagent/pkg/tools/builtin"
+	"github.com/docker/cagent/pkg/tools/mcp"
+)
+
+// ToolsetCreator is a function that creates a toolset based on the provided configuration
+type ToolsetCreator func(ctx context.Context, toolset latest.Toolset, parentDir string, envProvider environment.Provider, runtimeConfig config.RuntimeConfig) (tools.ToolSet, error)
+
+// ToolsetRegistry manages the registration of toolset creators by type
+type ToolsetRegistry struct {
+	creators map[string]ToolsetCreator
+}
+
+// NewToolsetRegistry creates a new empty toolset registry
+func NewToolsetRegistry() *ToolsetRegistry {
+	return &ToolsetRegistry{
+		creators: make(map[string]ToolsetCreator),
+	}
+}
+
+// Register adds a new toolset creator for the given type
+func (r *ToolsetRegistry) Register(toolsetType string, creator ToolsetCreator) {
+	r.creators[toolsetType] = creator
+}
+
+// Get retrieves a toolset creator for the given type
+func (r *ToolsetRegistry) Get(toolsetType string) (ToolsetCreator, bool) {
+	creator, ok := r.creators[toolsetType]
+	return creator, ok
+}
+
+// CreateTool creates a toolset using the registered creator for the given type
+func (r *ToolsetRegistry) CreateTool(ctx context.Context, toolset latest.Toolset, parentDir string, envProvider environment.Provider, runtimeConfig config.RuntimeConfig) (tools.ToolSet, error) {
+	creator, ok := r.Get(toolset.Type)
+	if !ok {
+		return nil, fmt.Errorf("unknown toolset type: %s", toolset.Type)
+	}
+	return creator(ctx, toolset, parentDir, envProvider, runtimeConfig)
+}
+
+func NewDefaultToolsetRegistry() *ToolsetRegistry {
+	r := NewToolsetRegistry()
+	// Register all built-in toolset creators
+	r.Register("todo", createTodoTool)
+	r.Register("memory", createMemoryTool)
+	r.Register("think", createThinkTool)
+	r.Register("shell", createShellTool)
+	r.Register("script", createScriptTool)
+	r.Register("filesystem", createFilesystemTool)
+	r.Register("fetch", createFetchTool)
+	r.Register("mcp", createMCPTool)
+	r.Register("api", createAPITool)
+	return r
+}
+
+func createTodoTool(ctx context.Context, toolset latest.Toolset, parentDir string, envProvider environment.Provider, runtimeConfig config.RuntimeConfig) (tools.ToolSet, error) {
+	if toolset.Shared {
+		return builtin.NewSharedTodoTool(), nil
+	}
+	return builtin.NewTodoTool(), nil
+}
+
+func createMemoryTool(ctx context.Context, toolset latest.Toolset, parentDir string, envProvider environment.Provider, runtimeConfig config.RuntimeConfig) (tools.ToolSet, error) {
+	var memoryPath string
+	if filepath.IsAbs(toolset.Path) {
+		memoryPath = ""
+	} else if wd, err := os.Getwd(); err == nil {
+		memoryPath = wd
+	} else {
+		memoryPath = parentDir
+	}
+
+	validatedMemoryPath, err := path.ValidatePathInDirectory(toolset.Path, memoryPath)
+	if err != nil {
+		return nil, fmt.Errorf("invalid memory database path: %w", err)
+	}
+	if err := os.MkdirAll(filepath.Dir(validatedMemoryPath), 0o700); err != nil {
+		return nil, fmt.Errorf("failed to create memory database directory: %w", err)
+	}
+
+	db, err := sqlite.NewMemoryDatabase(validatedMemoryPath)
+	if err != nil {
+		return nil, fmt.Errorf("failed to create memory database: %w", err)
+	}
+
+	return builtin.NewMemoryTool(db), nil
+}
+
+func createThinkTool(ctx context.Context, toolset latest.Toolset, parentDir string, envProvider environment.Provider, runtimeConfig config.RuntimeConfig) (tools.ToolSet, error) {
+	return builtin.NewThinkTool(), nil
+}
+
+func createShellTool(ctx context.Context, toolset latest.Toolset, parentDir string, envProvider environment.Provider, runtimeConfig config.RuntimeConfig) (tools.ToolSet, error) {
+	env, err := environment.ExpandAll(ctx, environment.ToValues(toolset.Env), envProvider)
+	if err != nil {
+		return nil, fmt.Errorf("failed to expand the tool's environment variables: %w", err)
+	}
+	env = append(env, os.Environ()...)
+	return builtin.NewShellTool(env), nil
+}
+
+func createScriptTool(ctx context.Context, toolset latest.Toolset, parentDir string, envProvider environment.Provider, runtimeConfig config.RuntimeConfig) (tools.ToolSet, error) {
+	if len(toolset.Shell) == 0 {
+		return nil, fmt.Errorf("shell is required for script toolset")
+	}
+
+	env, err := environment.ExpandAll(ctx, environment.ToValues(toolset.Env), envProvider)
+	if err != nil {
+		return nil, fmt.Errorf("failed to expand the tool's environment variables: %w", err)
+	}
+	env = append(env, os.Environ()...)
+	return builtin.NewScriptShellTool(toolset.Shell, env), nil
+}
+
+func createFilesystemTool(ctx context.Context, toolset latest.Toolset, parentDir string, envProvider environment.Provider, runtimeConfig config.RuntimeConfig) (tools.ToolSet, error) {
+	wd := runtimeConfig.WorkingDir
+	if wd == "" {
+		var err error
+		wd, err = os.Getwd()
+		if err != nil {
+			return nil, fmt.Errorf("failed to get working directory: %w", err)
+		}
+	}
+
+	var opts []builtin.FileSystemOpt
+
+	// Handle ignore_vcs configuration (default to true)
+	ignoreVCS := true
+	if toolset.IgnoreVCS != nil {
+		ignoreVCS = *toolset.IgnoreVCS
+	}
+	opts = append(opts, builtin.WithIgnoreVCS(ignoreVCS))
+
+	// Handle post-edit commands
+	if len(toolset.PostEdit) > 0 {
+		postEditConfigs := make([]builtin.PostEditConfig, len(toolset.PostEdit))
+		for i, pe := range toolset.PostEdit {
+			postEditConfigs[i] = builtin.PostEditConfig{
+				Path: pe.Path,
+				Cmd:  pe.Cmd,
+			}
+		}
+		opts = append(opts, builtin.WithPostEditCommands(postEditConfigs))
+	}
+
+	return builtin.NewFilesystemTool([]string{wd}, opts...), nil
+}
+
+func createAPITool(ctx context.Context, toolset latest.Toolset, parentDir string, envProvider environment.Provider, runtimeConfig config.RuntimeConfig) (tools.ToolSet, error) {
+	if toolset.APIConfig.Endpoint == "" {
+		return nil, fmt.Errorf("api tool requires an endpoint in api_config")
+	}
+
+	toolset.APIConfig.Headers = js.Expand(ctx, toolset.APIConfig.Headers, envProvider)
+
+	return builtin.NewAPITool(toolset.APIConfig), nil
+}
+
+func createFetchTool(ctx context.Context, toolset latest.Toolset, parentDir string, envProvider environment.Provider, runtimeConfig config.RuntimeConfig) (tools.ToolSet, error) {
+	var opts []builtin.FetchToolOption
+	if toolset.Timeout > 0 {
+		timeout := time.Duration(toolset.Timeout) * time.Second
+		opts = append(opts, builtin.WithTimeout(timeout))
+	}
+	return builtin.NewFetchTool(opts...), nil
+}
+
+func createMCPTool(ctx context.Context, toolset latest.Toolset, parentDir string, envProvider environment.Provider, runtimeConfig config.RuntimeConfig) (tools.ToolSet, error) {
+	// MCP tool has three different modes: ref, command, and remote
+	if toolset.Ref != "" {
+		mcpServerName := gateway.ParseServerRef(toolset.Ref)
+		serverSpec, err := gateway.ServerSpec(ctx, mcpServerName)
+		if err != nil {
+			return nil, fmt.Errorf("fetching MCP server spec for %q: %w", mcpServerName, err)
+		}
+
+		// TODO(dga): until the MCP Gateway supports oauth with cagent, we fetch the remote url and directly connect to it.
+		if serverSpec.Type == "remote" {
+			return mcp.NewRemoteToolset(serverSpec.Remote.URL, serverSpec.Remote.TransportType, nil), nil
+		}
+
+		return mcp.NewGatewayToolset(ctx, mcpServerName, toolset.Config, envProvider)
+	}
+
+	if toolset.Command != "" {
+		env, err := environment.ExpandAll(ctx, environment.ToValues(toolset.Env), envProvider)
+		if err != nil {
+			return nil, fmt.Errorf("failed to expand the tool's environment variables: %w", err)
+		}
+		env = append(env, os.Environ()...)
+		return mcp.NewToolsetCommand(toolset.Command, toolset.Args, env), nil
+	}
+
+	if toolset.Remote.URL != "" {
+		headers := map[string]string{}
+		for k, v := range toolset.Remote.Headers {
+			expanded, err := environment.Expand(ctx, v, envProvider)
+			if err != nil {
+				return nil, fmt.Errorf("failed to expand header '%s': %w", k, err)
+			}
+
+			headers[k] = expanded
+		}
+
+		return mcp.NewRemoteToolset(toolset.Remote.URL, toolset.Remote.TransportType, headers), nil
+	}
+
+	return nil, fmt.Errorf("mcp toolset requires either ref, command, or remote configuration")
+}

pkg/teamloader/sources.go

@@ -0,0 +1,77 @@
+package teamloader
+
+import (
+	"fmt"
+	"os"
+	"path/filepath"
+)
+
+type AgentSource interface {
+	Name() string
+	ParentDir() string
+	Read() ([]byte, error)
+}
+
+// fileSource is used to load an agent configuration from a YAML file.
+type fileSource struct {
+	path string
+}
+
+func NewFileSource(path string) AgentSource {
+	return fileSource{
+		path: path,
+	}
+}
+
+func (a fileSource) Name() string {
+	return filepath.Base(a.path)
+}
+
+func (a fileSource) ParentDir() string {
+	return filepath.Dir(a.path)
+}
+
+func (a fileSource) Read() ([]byte, error) {
+	parentDir := a.ParentDir()
+	fs, err := os.OpenRoot(parentDir)
+	if err != nil {
+		return nil, fmt.Errorf("opening filesystem %s: %w", parentDir, err)
+	}
+
+	fileName := a.Name()
+	data, err := fs.ReadFile(fileName)
+	if err != nil {
+		return nil, fmt.Errorf("reading config file %s: %w", fileName, err)
+	}
+
+	return data, nil
+}
+
+// bytesSource is used to load an agent configuration from a []byte.
+type bytesSource struct {
+	workingDir string
+	data       []byte
+}
+
+func NewBytesSource(workingDir string, data []byte) AgentSource {
+	// TODO(dga): this is not perfect but ok for now
+	if workingDir == "" {
+		workingDir = "."
+	}
+	return bytesSource{
+		workingDir: workingDir,
+		data:       data,
+	}
+}
+
+func (a bytesSource) Name() string {
+	return ""
+}
+
+func (a bytesSource) ParentDir() string {
+	return a.workingDir
+}
+
+func (a bytesSource) Read() ([]byte, error) {
+	return a.data, nil
+}