Use the oci source in the server

Also cleanup

Signed-off-by: Djordje Lukic <djordje.lukic@docker.com>

Author: rumpl

cmd/root/api.go

@@ -4,18 +4,14 @@ import (
 	"fmt"
 	"log/slog"
 	"os"
-	"path/filepath"
-	"time"
 
 	"github.com/spf13/cobra"
 
 	"github.com/docker/cagent/pkg/agentfile"
 	"github.com/docker/cagent/pkg/cli"
 	"github.com/docker/cagent/pkg/config"
-	"github.com/docker/cagent/pkg/remote"
 	"github.com/docker/cagent/pkg/server"
 	"github.com/docker/cagent/pkg/session"
-	"github.com/docker/cagent/pkg/teamloader"
 	"github.com/docker/cagent/pkg/telemetry"
 )
 
@@ -73,96 +69,21 @@ func (f *apiFlags) runAPICommand(cmd *cobra.Command, args []string) error {
 
 	slog.Debug("Starting server", "agents", agentsPath, "addr", ln.Addr().String())
 
-	resolvedPath, err := agentfile.Resolve(ctx, out, agentsPath)
-	if err != nil {
-		return err
-	}
-
 	sessionStore, err := session.NewSQLiteSessionStore(f.sessionDB)
 	if err != nil {
 		return fmt.Errorf("failed to create session store: %w", err)
 	}
 
-	var opts []server.Opt
-
-	if !agentfile.IsOCIReference(agentsPath) {
-		stat, err := os.Stat(resolvedPath)
-		if err != nil {
-			return fmt.Errorf("failed to stat agents path: %w", err)
-		}
-		if stat.IsDir() {
-			// For directories: only set agentsDir, not agentsPath
-			opts = append(opts, server.WithAgentsDir(resolvedPath))
-		} else {
-			opts = append(opts, server.WithAgentsPath(resolvedPath), server.WithAgentsDir(filepath.Dir(resolvedPath)))
-		}
-	}
-
-	teams, err := teamloader.LoadTeams(ctx, resolvedPath, &f.runConfig)
+	sources, err := agentfile.ResolveSources(ctx, nil, agentsPath)
 	if err != nil {
-		return fmt.Errorf("failed to load teams: %w", err)
-	}
-
-	// For OCI refs: store the reference for later per-session reloading, then clean up temp file
-	if agentfile.IsOCIReference(agentsPath) {
-		teamKey := filepath.Base(resolvedPath)
-		opts = append(opts, server.WithOCIRef(teamKey, agentsPath))
-
-		if err := os.Remove(resolvedPath); err != nil {
-			slog.Warn("Failed to remove temporary OCI file", "path", resolvedPath, "error", err)
-		} else {
-			slog.Debug("Cleaned up temporary OCI file", "path", resolvedPath)
-		}
+		return fmt.Errorf("failed to resolve agent sources: %w", err)
 	}
-
-	defer func() {
-		for _, team := range teams {
-			if err := team.StopToolSets(ctx); err != nil {
-				slog.Error("Failed to stop tool sets", "error", err)
-			}
-		}
-	}()
-
-	s, err := server.New(sessionStore, &f.runConfig, teams, opts...)
+	s, err := server.New(sessionStore, &f.runConfig, sources)
 	if err != nil {
 		return fmt.Errorf("failed to create server: %w", err)
 	}
 
-	// Start background auto-pull for OCI references if enabled
-	if f.pullIntervalMins > 0 {
-		go func() {
-			ticker := time.NewTicker(time.Duration(f.pullIntervalMins) * time.Minute)
-			defer ticker.Stop()
-
-			slog.Info("Auto-pull enabled for OCI reference", "reference", agentsPath, "interval_minutes", f.pullIntervalMins)
-
-			for {
-				select {
-				case <-ctx.Done():
-					return
-				case <-ticker.C:
-					slog.Info("Auto-pulling OCI reference", "reference", agentsPath)
-					if _, err := remote.Pull(ctx, agentsPath, false); err != nil {
-						slog.Error("Failed to auto-pull OCI reference", "reference", agentsPath, "error", err)
-						continue
-					}
-
-					// Resolve the OCI reference to get the updated file path
-					newResolvedPath, err := agentfile.Resolve(ctx, out, agentsPath)
-					if err != nil {
-						slog.Error("Failed to resolve OCI reference after pull", "reference", agentsPath, "error", err)
-						continue
-					}
-
-					if err := s.ReloadTeams(ctx, newResolvedPath); err != nil {
-						slog.Error("Failed to reload teams", "reference", agentsPath, "error", err)
-					} else {
-						slog.Info("Successfully reloaded teams from updated OCI reference", "reference", agentsPath)
-					}
-				}
-			}
-		}()
-	}
+	// TODO(rumpl): implement pull interval
 
 	return s.Serve(ctx, ln)
 }

pkg/acp/agent.go

@@ -22,7 +22,7 @@ import (
 type Agent struct {
 	conn          *acp.AgentSideConnection
 	team          *team.Team
-	agentFilename string
+	source        teamloader.AgentSource
 	runtimeConfig *config.RuntimeConfig
 	sessions      map[string]*Session
 	mu            sync.Mutex
@@ -39,9 +39,9 @@ type Session struct {
 }
 
 // NewAgent creates a new ACP agent
-func NewAgent(agentFilename string, runtimeConfig *config.RuntimeConfig) *Agent {
+func NewAgent(source teamloader.AgentSource, runtimeConfig *config.RuntimeConfig) *Agent {
 	return &Agent{
-		agentFilename: agentFilename,
+		source:        source,
 		runtimeConfig: runtimeConfig,
 		sessions:      make(map[string]*Session),
 	}
@@ -69,7 +69,7 @@ func (a *Agent) Initialize(ctx context.Context, params acp.InitializeRequest) (a
 
 	a.mu.Lock()
 	defer a.mu.Unlock()
-	t, err := teamloader.Load(ctx, a.agentFilename, a.runtimeConfig, teamloader.WithToolsetRegistry(createToolsetRegistry(a)))
+	t, err := teamloader.LoadFrom(ctx, a.source, a.runtimeConfig, teamloader.WithToolsetRegistry(createToolsetRegistry(a)))
 	if err != nil {
 		return acp.InitializeResponse{}, fmt.Errorf("failed to load teams: %w", err)
 	}

pkg/acp/run.go

@@ -11,19 +11,15 @@ import (
 	"github.com/docker/cagent/pkg/config"
 )
 
-type discardOutput struct{}
-
-func (d *discardOutput) Printf(string, ...any) {}
-
 func Run(ctx context.Context, agentFilename string, stdin io.Reader, stdout io.Writer, runConfig *config.RuntimeConfig) error {
 	slog.Debug("Starting ACP server", "agent", agentFilename)
 
-	agentFilename, err := agentfile.Resolve(ctx, &discardOutput{}, agentFilename)
+	source, err := agentfile.ResolveSource(ctx, nil, agentFilename)
 	if err != nil {
 		return err
 	}
 
-	acpAgent := NewAgent(agentFilename, runConfig)
+	acpAgent := NewAgent(source, runConfig)
 	conn := acpsdk.NewAgentSideConnection(acpAgent, stdout, stdin)
 	conn.SetLogger(slog.Default())
 	acpAgent.SetAgentConnection(conn)

pkg/agentfile/resolver.go

@@ -3,6 +3,7 @@ package agentfile
 import (
 	"context"
 	_ "embed"
+	"fmt"
 	"log/slog"
 	"os"
 	"path/filepath"
@@ -11,6 +12,7 @@ import (
 	"github.com/google/go-containerregistry/pkg/name"
 
 	"github.com/docker/cagent/pkg/aliases"
+	"github.com/docker/cagent/pkg/reference"
 	"github.com/docker/cagent/pkg/teamloader"
 )
 
@@ -21,11 +23,56 @@ type Printer interface {
 	Printf(format string, a ...any)
 }
 
+// ResolveSource resolves an agent file reference (local file or OCI image) to a local file path
+// For OCI references, always checks remote for updates but falls back to local cache if offline
+func ResolveSources(ctx context.Context, out Printer, agentFilename string) (teamloader.AgentSources, error) {
+	resolvedPath, err := resolve(agentFilename)
+	if err != nil {
+		if IsOCIReference(agentFilename) {
+			return map[string]teamloader.AgentSource{reference.OciRefToFilename(agentFilename): teamloader.NewOCISource(agentFilename)}, nil
+		}
+		return nil, err
+	}
+
+	if resolvedPath == "default" {
+		return map[string]teamloader.AgentSource{"default": teamloader.NewBytesSource(defaultAgent)}, nil
+	}
+	if isLocalFile(resolvedPath) {
+		return map[string]teamloader.AgentSource{resolvedPath: teamloader.NewFileSource(resolvedPath)}, nil
+	}
+	if dirExists(resolvedPath) {
+		sources := make(teamloader.AgentSources)
+		entries, err := os.ReadDir(resolvedPath)
+		if err != nil {
+			return nil, fmt.Errorf("reading agents directory %s: %w", resolvedPath, err)
+		}
+		for _, entry := range entries {
+			if entry.IsDir() {
+				continue
+			}
+			ext := strings.ToLower(filepath.Ext(entry.Name()))
+			if ext != ".yaml" && ext != ".yml" {
+				continue
+			}
+			a := filepath.Join(resolvedPath, entry.Name())
+			sources[a], err = ResolveSource(ctx, out, a)
+			if err != nil {
+				return nil, err
+			}
+		}
+		return sources, nil
+	}
+	return map[string]teamloader.AgentSource{resolvedPath: teamloader.NewOCISource(agentFilename)}, nil
+}
+
 // ResolveSource resolves an agent file reference (local file or OCI image) to a local file path
 // For OCI references, always checks remote for updates but falls back to local cache if offline
 func ResolveSource(ctx context.Context, out Printer, agentFilename string) (teamloader.AgentSource, error) {
-	resolvedPath, err := Resolve(ctx, out, agentFilename)
+	resolvedPath, err := resolve(agentFilename)
 	if err != nil {
+		if IsOCIReference(agentFilename) {
+			return teamloader.NewOCISource(agentFilename), nil
+		}
 		return nil, err
 	}
 
@@ -35,11 +82,11 @@ func ResolveSource(ctx context.Context, out Printer, agentFilename string) (team
 	if isLocalFile(resolvedPath) {
 		return teamloader.NewFileSource(resolvedPath), nil
 	}
-	return teamloader.NewOCISource(resolvedPath), nil
+	return teamloader.NewOCISource(agentFilename), nil
 }
 
-// Resolve resolves an agent file reference (local file or OCI image) to a local file path
-func Resolve(ctx context.Context, out Printer, agentFilename string) (string, error) {
+// resolve resolves an agent reference, handling aliases and defaults
+func resolve(agentFilename string) (string, error) {
 	if agentFilename == "" {
 		agentFilename = "default"
 	}
@@ -57,13 +104,25 @@ func Resolve(ctx context.Context, out Printer, agentFilename string) (string, er
 		return "default", nil
 	}
 
-	return agentFilename, nil
+	abs, err := filepath.Abs(agentFilename)
+	if err != nil {
+		return "", err
+	}
+
+	return abs, nil
 }
 
 // fileExists checks if a file exists at the given path
 func fileExists(path string) bool {
-	_, err := os.Stat(path)
-	exists := err == nil
+	s, err := os.Stat(path)
+	exists := err == nil && !s.IsDir()
+	return exists
+}
+
+// fileExists checks if a file exists at the given path
+func dirExists(path string) bool {
+	s, err := os.Stat(path)
+	exists := err == nil && s.IsDir()
 	return exists
 }