From bd2414dc25b8bae3615ee7f9d07b345400b1690f Mon Sep 17 00:00:00 2001 From: David Gageot Date: Sun, 22 Mar 2026 12:42:53 +0100 Subject: [PATCH] Make run.go easier to read Assisted-By: docker-agent --- cmd/root/record.go | 61 +++++++------------------------- cmd/root/run.go | 56 +++++++----------------------- pkg/profiling/profiling.go | 65 ++++++++++++++++++++++++++++++++++ pkg/recording/recording.go | 71 ++++++++++++++++++++++++++++++++++++++ 4 files changed, 161 insertions(+), 92 deletions(-) create mode 100644 pkg/profiling/profiling.go create mode 100644 pkg/recording/recording.go diff --git a/cmd/root/record.go b/cmd/root/record.go index 8b4e9efe6..11ca7f62c 100644 --- a/cmd/root/record.go +++ b/cmd/root/record.go @@ -1,71 +1,36 @@ package root import ( - "fmt" - "log/slog" - "strings" - "time" - "github.com/docker/docker-agent/pkg/config" - "github.com/docker/docker-agent/pkg/fake" + "github.com/docker/docker-agent/pkg/recording" ) // setupFakeProxy starts a fake proxy if fakeResponses is non-empty. -// streamDelayMs controls simulated streaming: 0 = disabled, >0 = delay in milliseconds between chunks. -// It returns a cleanup function that must be called when done (typically via defer). +// It configures the runtime config's ModelsGateway to point to the proxy. func setupFakeProxy(fakeResponses string, streamDelayMs int, runConfig *config.RuntimeConfig) (cleanup func() error, err error) { - if fakeResponses == "" { - return func() error { return nil }, nil - } - - // Normalize path by stripping .yaml suffix (go-vcr adds it automatically) - fakeResponses = strings.TrimSuffix(fakeResponses, ".yaml") - - var opts []fake.ProxyOption - if streamDelayMs > 0 { - opts = append(opts, - fake.WithSimulateStream(true), - fake.WithStreamChunkDelay(time.Duration(streamDelayMs)*time.Millisecond), - ) - } - - proxyURL, cleanupFn, err := fake.StartProxy(fakeResponses, opts...) + proxyURL, cleanupFn, err := recording.SetupFakeProxy(fakeResponses, streamDelayMs) if err != nil { - return nil, fmt.Errorf("failed to start fake proxy: %w", err) + return nil, err } - runConfig.ModelsGateway = proxyURL - slog.Info("Fake mode enabled", "cassette", fakeResponses, "proxy", proxyURL) + if proxyURL != "" { + runConfig.ModelsGateway = proxyURL + } return cleanupFn, nil } // setupRecordingProxy starts a recording proxy if recordPath is non-empty. -// It handles auto-generating a filename when recordPath is "true" (from NoOptDefVal), -// and normalizes the path by stripping any .yaml suffix. -// Returns the cassette path (with .yaml extension) and a cleanup function. -// The cleanup function must be called when done (typically via defer). +// It configures the runtime config's ModelsGateway to point to the proxy. func setupRecordingProxy(recordPath string, runConfig *config.RuntimeConfig) (cassettePath string, cleanup func() error, err error) { - if recordPath == "" { - return "", func() error { return nil }, nil - } - - // Handle auto-generated filename (from NoOptDefVal) - if recordPath == "true" { - recordPath = fmt.Sprintf("cagent-recording-%d", time.Now().Unix()) - } else { - recordPath = strings.TrimSuffix(recordPath, ".yaml") - } - - proxyURL, cleanupFn, err := fake.StartRecordingProxy(recordPath) + cassettePath, proxyURL, cleanupFn, err := recording.SetupRecordingProxy(recordPath) if err != nil { - return "", nil, fmt.Errorf("failed to start recording proxy: %w", err) + return "", nil, err } - runConfig.ModelsGateway = proxyURL - cassettePath = recordPath + ".yaml" - - slog.Info("Recording mode enabled", "cassette", cassettePath, "proxy", proxyURL) + if proxyURL != "" { + runConfig.ModelsGateway = proxyURL + } return cassettePath, cleanupFn, nil } diff --git a/cmd/root/run.go b/cmd/root/run.go index 6070b1c3f..c650c98f6 100644 --- a/cmd/root/run.go +++ b/cmd/root/run.go @@ -8,8 +8,6 @@ import ( "log/slog" "os" "path/filepath" - goruntime "runtime" - "runtime/pprof" "sync" "time" @@ -22,6 +20,7 @@ import ( "github.com/docker/docker-agent/pkg/cli" "github.com/docker/docker-agent/pkg/config" "github.com/docker/docker-agent/pkg/paths" + "github.com/docker/docker-agent/pkg/profiling" "github.com/docker/docker-agent/pkg/runtime" "github.com/docker/docker-agent/pkg/session" "github.com/docker/docker-agent/pkg/sessiontitle" @@ -145,37 +144,16 @@ func (f *runExecFlags) runRunCommand(cmd *cobra.Command, args []string) error { func (f *runExecFlags) runOrExec(ctx context.Context, out *cli.Printer, args []string, useTUI bool) error { slog.Debug("Starting agent", "agent", f.agentName) - // Start CPU profiling if requested - if f.cpuProfile != "" { - pf, err := os.Create(f.cpuProfile) - if err != nil { - return fmt.Errorf("failed to create CPU profile: %w", err) - } - if err := pprof.StartCPUProfile(pf); err != nil { - pf.Close() - return fmt.Errorf("failed to start CPU profile: %w", err) - } - defer pprof.StopCPUProfile() - defer pf.Close() - slog.Info("CPU profiling enabled", "file", f.cpuProfile) - } - - // Write memory profile at exit if requested - if f.memProfile != "" { - defer func() { - mf, err := os.Create(f.memProfile) - if err != nil { - slog.Error("Failed to create memory profile", "error", err) - return - } - defer mf.Close() - goruntime.GC() // Get up-to-date statistics - if err := pprof.WriteHeapProfile(mf); err != nil { - slog.Error("Failed to write memory profile", "error", err) - } - slog.Info("Memory profile written", "file", f.memProfile) - }() + // Start profiling if requested + stopProfiling, err := profiling.Start(f.cpuProfile, f.memProfile) + if err != nil { + return err } + defer func() { + if err := stopProfiling(); err != nil { + slog.Error("Profiling cleanup failed", "error", err) + } + }() var agentFileName string if len(args) > 0 { @@ -271,10 +249,6 @@ func (f *runExecFlags) runOrExec(ctx context.Context, out *cli.Printer, args []s } defer initialTeamCleanup() - if useTUI { - applyTheme() - } - if f.dryRun { out.Println("Dry run mode enabled. Agent initialized but will not execute.") return nil @@ -284,19 +258,13 @@ func (f *runExecFlags) runOrExec(ctx context.Context, out *cli.Printer, args []s return f.handleExecMode(ctx, out, rt, sess, args) } + applyTheme() opts, err := f.buildAppOpts(args) if err != nil { return err } - var sessStore session.Store - switch typedRt := rt.(type) { - case *runtime.LocalRuntime: - sessStore = typedRt.SessionStore() - case *runtime.PersistentRuntime: - sessStore = typedRt.SessionStore() - } - + sessStore := rt.SessionStore() return runTUI(ctx, rt, sess, f.createSessionSpawner(agentSource, sessStore), initialTeamCleanup, opts...) } diff --git a/pkg/profiling/profiling.go b/pkg/profiling/profiling.go new file mode 100644 index 000000000..5a1248a28 --- /dev/null +++ b/pkg/profiling/profiling.go @@ -0,0 +1,65 @@ +// Package profiling provides helpers for CPU and memory profiling. +package profiling + +import ( + "errors" + "fmt" + "os" + "runtime" + "runtime/pprof" +) + +// Stop is a function returned by Start that stops profiling and flushes +// any buffered data. It must be called (typically via defer) when the +// profiled section of code completes. +type Stop func() error + +// Start begins CPU and/or memory profiling based on the provided file +// paths. Pass an empty string to skip the corresponding profile. +// The returned Stop function must be called to finalise the profiles. +func Start(cpuProfile, memProfile string) (Stop, error) { + var closers []func() error + + if cpuProfile != "" { + f, err := os.Create(cpuProfile) + if err != nil { + return noop, fmt.Errorf("failed to create CPU profile: %w", err) + } + if err := pprof.StartCPUProfile(f); err != nil { + f.Close() + return noop, fmt.Errorf("failed to start CPU profile: %w", err) + } + closers = append(closers, func() error { + pprof.StopCPUProfile() + return f.Close() + }) + } + + if memProfile != "" { + closers = append(closers, func() error { + f, err := os.Create(memProfile) + if err != nil { + return fmt.Errorf("failed to create memory profile: %w", err) + } + defer f.Close() + runtime.GC() + if err := pprof.WriteHeapProfile(f); err != nil { + return fmt.Errorf("failed to write memory profile: %w", err) + } + return nil + }) + } + + return func() error { + // Run in reverse order so CPU profile is stopped before mem profile is written. + var errs []error + for i := len(closers) - 1; i >= 0; i-- { + if err := closers[i](); err != nil { + errs = append(errs, err) + } + } + return errors.Join(errs...) + }, nil +} + +func noop() error { return nil } diff --git a/pkg/recording/recording.go b/pkg/recording/recording.go new file mode 100644 index 000000000..474f2195f --- /dev/null +++ b/pkg/recording/recording.go @@ -0,0 +1,71 @@ +// Package recording provides helpers for recording and replaying AI API interactions. +package recording + +import ( + "fmt" + "log/slog" + "strings" + "time" + + "github.com/docker/docker-agent/pkg/fake" +) + +// SetupFakeProxy starts a fake proxy if fakeResponses is non-empty. +// streamDelayMs controls simulated streaming: 0 = disabled, >0 = delay in milliseconds between chunks. +// It returns the proxy URL and a cleanup function that must be called when done (typically via defer). +func SetupFakeProxy(fakeResponses string, streamDelayMs int) (proxyURL string, cleanup func() error, err error) { + if fakeResponses == "" { + return "", noop, nil + } + + // Normalize path by stripping .yaml suffix (go-vcr adds it automatically) + fakeResponses = strings.TrimSuffix(fakeResponses, ".yaml") + + var opts []fake.ProxyOption + if streamDelayMs > 0 { + opts = append(opts, + fake.WithSimulateStream(true), + fake.WithStreamChunkDelay(time.Duration(streamDelayMs)*time.Millisecond), + ) + } + + proxyURL, cleanupFn, err := fake.StartProxy(fakeResponses, opts...) + if err != nil { + return "", nil, fmt.Errorf("failed to start fake proxy: %w", err) + } + + slog.Info("Fake mode enabled", "cassette", fakeResponses, "proxy", proxyURL) + + return proxyURL, cleanupFn, nil +} + +// SetupRecordingProxy starts a recording proxy if recordPath is non-empty. +// It handles auto-generating a filename when recordPath is "true" (from NoOptDefVal), +// and normalizes the path by stripping any .yaml suffix. +// Returns the cassette path (with .yaml extension), the proxy URL, and a cleanup function. +// The cleanup function must be called when done (typically via defer). +func SetupRecordingProxy(recordPath string) (cassettePath, proxyURL string, cleanup func() error, err error) { + if recordPath == "" { + return "", "", noop, nil + } + + // Handle auto-generated filename (from NoOptDefVal) + if recordPath == "true" { + recordPath = fmt.Sprintf("cagent-recording-%d", time.Now().Unix()) + } else { + recordPath = strings.TrimSuffix(recordPath, ".yaml") + } + + proxyURL, cleanupFn, err := fake.StartRecordingProxy(recordPath) + if err != nil { + return "", "", nil, fmt.Errorf("failed to start recording proxy: %w", err) + } + + cassettePath = recordPath + ".yaml" + + slog.Info("Recording mode enabled", "cassette", cassettePath, "proxy", proxyURL) + + return cassettePath, proxyURL, cleanupFn, nil +} + +func noop() error { return nil }