Merge pull request #1734 from dgageot/cmdroot

Small improvements to cmd/root

Author: rumpl

cmd/root/a2a.go

@@ -1,14 +1,11 @@
 package root
 
 import (
-	"fmt"
-
 	"github.com/spf13/cobra"
 
 	"github.com/docker/cagent/pkg/a2a"
 	"github.com/docker/cagent/pkg/cli"
 	"github.com/docker/cagent/pkg/config"
-	"github.com/docker/cagent/pkg/server"
 	"github.com/docker/cagent/pkg/telemetry"
 )
 
@@ -46,15 +43,10 @@ func (f *a2aFlags) runA2ACommand(cmd *cobra.Command, args []string) error {
 	out := cli.NewPrinter(cmd.OutOrStdout())
 	agentFilename := args[0]
 
-	// Listen as early as possible
-	ln, err := server.Listen(ctx, f.listenAddr)
+	ln, err := listenAndCloseOnCancel(ctx, f.listenAddr)
 	if err != nil {
-		return fmt.Errorf("failed to listen on %s: %w", f.listenAddr, err)
+		return err
 	}
-	go func() {
-		<-ctx.Done()
-		_ = ln.Close()
-	}()
 
 	out.Println("Listening on", ln.Addr().String())
 	return a2a.Run(ctx, agentFilename, f.agentName, &f.runConfig, ln)

cmd/root/api.go

@@ -123,7 +123,7 @@ func (f *apiFlags) runAPICommand(cmd *cobra.Command, args []string) error {
 	// Start recording proxy if --record is specified
 	if _, recordCleanup, err := setupRecordingProxy(f.recordPath, &f.runConfig); err != nil {
 		return err
-	} else if recordCleanup != nil {
+	} else {
 		defer func() {
 			if err := recordCleanup(); err != nil {
 				slog.Error("Failed to cleanup recording proxy", "error", err)
@@ -135,14 +135,10 @@ func (f *apiFlags) runAPICommand(cmd *cobra.Command, args []string) error {
 		return fmt.Errorf("--pull-interval flag can only be used with OCI references, not local files")
 	}
 
-	ln, err := server.Listen(ctx, f.listenAddr)
+	ln, err := listenAndCloseOnCancel(ctx, f.listenAddr)
 	if err != nil {
-		return fmt.Errorf("failed to listen on %s: %w", f.listenAddr, err)
+		return err
 	}
-	go func() {
-		<-ctx.Done()
-		_ = ln.Close()
-	}()
 
 	out.Println("Listening on", ln.Addr().String())
 

cmd/root/catalog.go

@@ -40,10 +40,8 @@ func newCatalogListCmd() *cobra.Command {
 func runCatalogListCommand(cmd *cobra.Command, args []string) error {
 	telemetry.TrackCommand("catalog", append([]string{"list"}, args...))
 
-	var org string
-	if len(args) == 0 {
-		org = "agentcatalog"
-	} else {
+	org := "agentcatalog"
+	if len(args) > 0 {
 		org = args[0]
 	}
 

cmd/root/debug.go

@@ -56,15 +56,10 @@ func newDebugCmd() *cobra.Command {
 	return cmd
 }
 
-// resolveSource resolves an agent file reference to a config source.
-func (f *debugFlags) resolveSource(agentFilename string) (config.Source, error) {
-	return config.Resolve(agentFilename, f.runConfig.EnvProvider())
-}
-
 // loadTeam loads an agent team from the given agent file and returns
 // a cleanup function that must be deferred by the caller.
 func (f *debugFlags) loadTeam(ctx context.Context, agentFilename string, opts ...teamloader.Opt) (*team.Team, func(), error) {
-	agentSource, err := f.resolveSource(agentFilename)
+	agentSource, err := config.Resolve(agentFilename, f.runConfig.EnvProvider())
 	if err != nil {
 		return nil, nil, err
 	}
@@ -86,7 +81,7 @@ func (f *debugFlags) loadTeam(ctx context.Context, agentFilename string, opts ..
 func (f *debugFlags) runDebugConfigCommand(cmd *cobra.Command, args []string) error {
 	telemetry.TrackCommand("debug", append([]string{"config"}, args...))
 
-	agentSource, err := f.resolveSource(args[0])
+	agentSource, err := config.Resolve(args[0], f.runConfig.EnvProvider())
 	if err != nil {
 		return err
 	}

cmd/root/exec.go

@@ -35,9 +35,7 @@ func newExecCmd() *cobra.Command {
 func (f *runExecFlags) runExecCommand(cmd *cobra.Command, args []string) error {
 	telemetry.TrackCommand("exec", args)
 
-	ctx := cmd.Context()
 	out := cli.NewPrinter(cmd.OutOrStdout())
 
-	tui := false
-	return f.runOrExec(ctx, out, args, tui)
+	return f.runOrExec(cmd.Context(), out, args, false)
 }

cmd/root/flags.go

@@ -1,8 +1,10 @@
 package root
 
 import (
+	"context"
 	"fmt"
 	"log/slog"
+	"net"
 	"os"
 	"path/filepath"
 	"strings"
@@ -11,6 +13,7 @@ import (
 
 	"github.com/docker/cagent/pkg/config"
 	"github.com/docker/cagent/pkg/config/latest"
+	"github.com/docker/cagent/pkg/server"
 	"github.com/docker/cagent/pkg/userconfig"
 )
 
@@ -113,3 +116,17 @@ func parseModelShorthand(s string) *latest.ModelConfig {
 	}
 	return nil
 }
+
+// listenAndCloseOnCancel starts a listener and spawns a goroutine
+// that closes it when the context is cancelled.
+func listenAndCloseOnCancel(ctx context.Context, addr string) (net.Listener, error) {
+	ln, err := server.Listen(ctx, addr)
+	if err != nil {
+		return nil, fmt.Errorf("failed to listen on %s: %w", addr, err)
+	}
+	go func() {
+		<-ctx.Done()
+		_ = ln.Close()
+	}()
+	return ln, nil
+}

cmd/root/mcp.go

@@ -1,13 +1,10 @@
 package root
 
 import (
-	"fmt"
-
 	"github.com/spf13/cobra"
 
 	"github.com/docker/cagent/pkg/config"
 	"github.com/docker/cagent/pkg/mcp"
-	"github.com/docker/cagent/pkg/server"
 	"github.com/docker/cagent/pkg/telemetry"
 )
 
@@ -52,15 +49,10 @@ func (f *mcpFlags) runMCPCommand(cmd *cobra.Command, args []string) error {
 		return mcp.StartMCPServer(ctx, agentFilename, f.agentName, &f.runConfig)
 	}
 
-	// Listen as early as possible
-	ln, err := server.Listen(ctx, f.listenAddr)
+	ln, err := listenAndCloseOnCancel(ctx, f.listenAddr)
 	if err != nil {
-		return fmt.Errorf("failed to listen on %s: %w", f.listenAddr, err)
+		return err
 	}
-	go func() {
-		<-ctx.Done()
-		_ = ln.Close()
-	}()
 
 	return mcp.StartHTTPServer(ctx, agentFilename, f.agentName, &f.runConfig, ln)
 }

cmd/root/otel.go

@@ -44,8 +44,9 @@ func initOTelSDK(ctx context.Context) (err error) {
 	}
 
 	// Configure tracer provider
-	var tracerProviderOpts []trace.TracerProviderOption
-	tracerProviderOpts = append(tracerProviderOpts, trace.WithResource(res))
+	tracerProviderOpts := []trace.TracerProviderOption{
+		trace.WithResource(res),
+	}
 
 	if traceExporter != nil {
 		tracerProviderOpts = append(tracerProviderOpts,

cmd/root/pull.go

@@ -6,7 +6,6 @@ import (
 	"os"
 	"strings"
 
-	"github.com/google/go-containerregistry/pkg/crane"
 	"github.com/spf13/cobra"
 
 	"github.com/docker/cagent/pkg/cli"
@@ -46,8 +45,7 @@ func (f *pullFlags) runPullCommand(cmd *cobra.Command, args []string) error {
 
 	out.Println("Pulling agent", registryRef)
 
-	var opts []crane.Option
-	_, err := remote.Pull(ctx, registryRef, f.force, opts...)
+	_, err := remote.Pull(ctx, registryRef, f.force)
 	if err != nil {
 		return fmt.Errorf("failed to pull artifact: %w", err)
 	}

cmd/root/run.go

@@ -234,7 +234,14 @@ func (f *runExecFlags) runOrExec(ctx context.Context, out *cli.Printer, args []s
 			return err
 		}
 
-		loadResult, err := f.loadAgentFrom(ctx, agentSource)
+		opts := []teamloader.Opt{
+			teamloader.WithModelOverrides(f.modelOverrides),
+		}
+		if len(f.promptFiles) > 0 {
+			opts = append(opts, teamloader.WithPromptFiles(f.promptFiles))
+		}
+
+		loadResult, err := teamloader.LoadWithConfig(ctx, agentSource, &f.runConfig, opts...)
 		if err != nil {
 			return err
 		}
@@ -275,58 +282,16 @@ func (f *runExecFlags) runOrExec(ctx context.Context, out *cli.Printer, args []s
 	return f.handleRunMode(ctx, rt, sess, args)
 }
 
-func (f *runExecFlags) loadAgentFrom(ctx context.Context, agentSource config.Source) (*teamloader.LoadResult, error) {
-	opts := []teamloader.Opt{
-		teamloader.WithModelOverrides(f.modelOverrides),
-	}
-	if len(f.promptFiles) > 0 {
-		opts = append(opts, teamloader.WithPromptFiles(f.promptFiles))
-	}
-
-	result, err := teamloader.LoadWithConfig(ctx, agentSource, &f.runConfig, opts...)
-	if err != nil {
-		return nil, err
-	}
-
-	return result, nil
-}
-
 func (f *runExecFlags) createRemoteRuntimeAndSession(ctx context.Context, originalFilename string) (runtime.Runtime, *session.Session, error) {
-	if f.connectRPC {
-		return f.createConnectRPCRuntimeAndSession(ctx, originalFilename)
-	}
-	return f.createHTTPRuntimeAndSession(ctx, originalFilename)
-}
-
-func (f *runExecFlags) createConnectRPCRuntimeAndSession(ctx context.Context, originalFilename string) (runtime.Runtime, *session.Session, error) {
-	connectClient, err := runtime.NewConnectRPCClient(f.remoteAddress)
-	if err != nil {
-		return nil, nil, fmt.Errorf("failed to create connect-rpc client: %w", err)
-	}
-
-	sessTemplate := session.New(
-		session.WithToolsApproved(f.autoApprove),
-	)
-
-	sess, err := connectClient.CreateSession(ctx, sessTemplate)
-	if err != nil {
-		return nil, nil, err
-	}
-
-	remoteRt, err := runtime.NewRemoteRuntime(connectClient,
-		runtime.WithRemoteCurrentAgent(f.agentName),
-		runtime.WithRemoteAgentFilename(originalFilename),
+	var (
+		client runtime.RemoteClient
+		err    error
 	)
-	if err != nil {
-		return nil, nil, fmt.Errorf("failed to create connect-rpc remote runtime: %w", err)
+	if f.connectRPC {
+		client, err = runtime.NewConnectRPCClient(f.remoteAddress)
+	} else {
+		client, err = runtime.NewClient(f.remoteAddress)
 	}
-
-	slog.Debug("Using connect-rpc remote runtime", "address", f.remoteAddress, "agent", f.agentName)
-	return remoteRt, sess, nil
-}
-
-func (f *runExecFlags) createHTTPRuntimeAndSession(ctx context.Context, originalFilename string) (runtime.Runtime, *session.Session, error) {
-	remoteClient, err := runtime.NewClient(f.remoteAddress)
 	if err != nil {
 		return nil, nil, fmt.Errorf("failed to create remote client: %w", err)
 	}
@@ -335,20 +300,20 @@ func (f *runExecFlags) createHTTPRuntimeAndSession(ctx context.Context, original
 		session.WithToolsApproved(f.autoApprove),
 	)
 
-	sess, err := remoteClient.CreateSession(ctx, sessTemplate)
+	sess, err := client.CreateSession(ctx, sessTemplate)
 	if err != nil {
 		return nil, nil, err
 	}
 
-	remoteRt, err := runtime.NewRemoteRuntime(remoteClient,
+	remoteRt, err := runtime.NewRemoteRuntime(client,
 		runtime.WithRemoteCurrentAgent(f.agentName),
 		runtime.WithRemoteAgentFilename(originalFilename),
 	)
 	if err != nil {
 		return nil, nil, fmt.Errorf("failed to create remote runtime: %w", err)
 	}
 
-	slog.Debug("Using remote runtime", "address", f.remoteAddress, "agent", f.agentName)
+	slog.Debug("Using remote runtime", "address", f.remoteAddress, "agent", f.agentName, "connect_rpc", f.connectRPC)
 	return remoteRt, sess, nil
 }