CLI Improvements (#115)

* CLI Improvements

* (feat): add bubble tea CUI

* add debug logging

* debug logging and CLI changes

* fix: no terminal in ci/cd

* add testing for new cli

* added testing for TUI

* Make File Exclusions be a flag + Testing

Author: mirackara

README.md

@@ -19,13 +19,13 @@ This project, its code, and the UX are under heavy development, and should be ex
 
 This tool doesn't interfere with your application's operation, and it doesn't make any changes to your code directly. Here's what happens when you run the tool:
 
-* It analyzes your code and suggests changes that allow the Go agent to capture telemetry data. 
-* You review the changes in the .diff file and decide which changes to add to your source code.
+* It analyzes your code and suggests changes that allow the Go agent to capture telemetry data.
+* You watch the progress in real-time as the tool scans and instruments your source files.
+* A diff file is generated with the proposed changes.
+* You apply the changes to your source code using `git apply <diff-file>`.
 
 As part of the analysis, this tool may invoke `go get` or other Go language toolchain commands which may modify your `go.mod` file, but not your actual source code.
 
-**IMPORTANT:** This tool can't detect if you already have New Relic instrumentation. Please only use this on applications without any instrumentation.
-
 ## What is instrumented?
 
 The scope of what this tool can instrument in your application is limited to these actions:
@@ -68,6 +68,10 @@ Installation Steps have been moved to: https://docs.newrelic.com/docs/apm/agents
     ```
     go run . instrument $MY_APP
     ```
+5.  Once the instrumentation is complete, the tool will output the path to a diff file. Apply the changes to your source code:
+    ```sh
+    git apply <path-to-diff-file>
+    ```
 
 
 ### Updating an existing version of Go Easy Instrumentation

cmd/instrument.go

@@ -5,11 +5,16 @@ import (
 	"fmt"
 	"os"
 	"path/filepath"
+	"strings"
 
+	"github.com/charmbracelet/bubbles/spinner"
+	tea "github.com/charmbracelet/bubbletea"
+	"github.com/charmbracelet/lipgloss"
 	"github.com/dave/dst/decorator"
 	"github.com/newrelic/go-easy-instrumentation/internal/comment"
 	"github.com/newrelic/go-easy-instrumentation/parser"
 	"github.com/spf13/cobra"
+	"golang.org/x/term"
 	"golang.org/x/tools/go/packages"
 )
 
@@ -20,12 +25,11 @@ const (
 	defaultAppName           = ""
 	defaultOutputFilePath    = ""
 	defaultDiffFileName      = "new-relic-instrumentation.diff"
-	defaultDebug             = false
 )
 
 var (
-	debug    bool
-	diffFile string
+	diffFile    string
+	excludeDirs string
 )
 
 var instrumentCmd = &cobra.Command{
@@ -73,56 +77,277 @@ func setOutputFilePath(outputFilePath, applicationPath string) (string, error) {
 
 const LoadMode = packages.LoadSyntax | packages.NeedForTest
 
-func Instrument(packagePath string) {
+// Bubble Tea Model
+type model struct {
+	spinner     spinner.Model
+	stepDesc    string
+	totalSteps  int
+	currentStep int
+	done        bool
+	err         error
+	packages    []*decorator.Package
+	pkgPath     string
+	sub         chan tea.Msg
+	outputFile  string
+}
+
+// Messages
+// Messages
+type progressMsg struct {
+	desc string
+}
+type pkgLoadedMsg []*decorator.Package
+type errMsg error
+type completedMsg struct{}
+
+func initialModel(pkgPath, outputFile string) model {
+	s := spinner.New()
+	s.Spinner = spinner.Dot
+	s.Style = lipgloss.NewStyle().Foreground(lipgloss.Color("#1CE783"))
+	return model{
+		spinner:    s,
+		stepDesc:   "Loading packages...",
+		totalSteps: 8,
+		pkgPath:    pkgPath,
+		outputFile: outputFile,
+	}
+}
+
+func (m model) Init() tea.Cmd {
+	return tea.Batch(
+		m.spinner.Tick,
+		waitForNext(m.sub),
+	)
+}
+
+func Instrument(packagePath string, patterns ...string) {
 	if packagePath == "" {
 		cobra.CheckErr("path argument cannot be empty")
 	}
-
 	if _, err := os.Stat(packagePath); err != nil {
 		cobra.CheckErr(fmt.Errorf("path argument \"%s\" is invalid: %v", packagePath, err))
 	}
-
 	outputFile, err := setOutputFilePath(diffFile, packagePath)
 	cobra.CheckErr(err)
 	if debug {
 		comment.EnableConsolePrinter(packagePath)
 	}
 
-	pkgs, err := decorator.Load(&packages.Config{Dir: packagePath, Mode: LoadMode, Tests: true}, defaultPackageName)
-	cobra.CheckErr(err)
+	// If debug mode is enabled or no terminal is available (CI/CD), run in text mode (no TUI)
+	if debug || !term.IsTerminal(int(os.Stdout.Fd())) {
+		runTextMode(packagePath, patterns, outputFile)
+		return
+	}
+
+	// Normal TUI mode
+	runTUIMode(packagePath, patterns, outputFile)
+}
+
+// runTextMode runs the instrumentation pipeline with plain text output to stdout.
+// This is used when the TUI is unavailable (e.g. CI/CD, piped output) or when
+// the --debug flag is enabled. It delegates to instrumentPackages for the core
+// logic and handles printing status and exit on error.
+func runTextMode(packagePath string, patterns []string, outputFile string) {
+	fmt.Printf("Instrumentation started for %s\n", packagePath)
+	fmt.Printf("Output file: %s\n\n", outputFile)
+
+	if err := instrumentPackages(packagePath, patterns, outputFile); err != nil {
+		fmt.Printf("Error: %v\n", err)
+		os.Exit(1)
+	}
+
+	fmt.Printf("\nDone! Changes written to: %s\nTip: Apply these changes with: git apply %s\n", outputFile, outputFile)
+}
+
+// instrumentPackages loads Go packages from packagePath, runs the full
+// instrumentation pipeline, and writes the resulting diff to outputFile.
+// It returns an error rather than exiting the process, making it safe to
+// call from both runTextMode (which handles os.Exit) and from tests.
+func instrumentPackages(packagePath string, patterns []string, outputFile string) error {
+	loadPatterns := patterns
+	if len(loadPatterns) == 0 {
+		loadPatterns = []string{defaultPackageName}
+	}
+
+	fmt.Println(" -> Loading packages...")
+	pkgs, err := decorator.Load(&packages.Config{Dir: packagePath, Mode: LoadMode, Tests: true}, loadPatterns...)
+	if err != nil {
+		return fmt.Errorf("loading packages: %w", err)
+	}
 
 	manager := parser.NewInstrumentationManager(pkgs, defaultAppName, defaultAgentVariableName, outputFile, packagePath)
-	err = manager.CreateDiffFile()
-	cobra.CheckErr(err)
 
-	err = manager.DetectDependencyIntegrations()
-	cobra.CheckErr(err)
+	steps := []struct {
+		desc string
+		fn   func() error
+	}{
+		{"Creating diff file", manager.CreateDiffFile},
+		{"Detecting dependencies", manager.DetectDependencyIntegrations},
+		{"Tracing package calls", manager.TracePackageCalls},
+		{"Scanning application", manager.ScanApplication},
+		{"Instrumenting application", manager.InstrumentApplication},
+		{"Resolving unit tests", manager.ResolveUnitTests},
+		{"Adding required modules", manager.AddRequiredModules},
+		{"Writing diff file", func() error {
+			comment.WriteAll()
+			return manager.WriteDiff(func(msg string) {})
+		}},
+	}
 
-	err = manager.TracePackageCalls()
-	cobra.CheckErr(err)
+	for _, step := range steps {
+		if err := step.fn(); err != nil {
+			return fmt.Errorf("%s: %w", step.desc, err)
+		}
+	}
 
-	err = manager.ScanApplication()
-	cobra.CheckErr(err)
+	return nil
+}
 
-	err = manager.InstrumentApplication()
-	cobra.CheckErr(err)
+func runTUIMode(packagePath string, patterns []string, outputFile string) {
+	// Channel to receive updates from the worker
+	updates := make(chan tea.Msg)
 
-	err = manager.ResolveUnitTests()
-	cobra.CheckErr(err)
+	// Worker goroutine
+	go func() {
+		loadPatterns := patterns
+		if len(loadPatterns) == 0 {
+			loadPatterns = []string{defaultPackageName}
+		}
 
-	err = manager.AddRequiredModules()
-	cobra.CheckErr(err)
+		pkgs, err := decorator.Load(&packages.Config{Dir: packagePath, Mode: LoadMode, Tests: true}, loadPatterns...)
+		if err != nil {
+			updates <- errMsg(err)
+			return
+		}
 
-	// write debug comments before writing diff so that
-	// diff file console log is still easy to see
-	comment.WriteAll()
-	err = manager.WriteDiff()
-	cobra.CheckErr(err)
+		updates <- pkgLoadedMsg(pkgs)
+
+		manager := parser.NewInstrumentationManager(pkgs, defaultAppName, defaultAgentVariableName, outputFile, packagePath)
+
+		steps := []struct {
+			desc string
+			fn   func() error
+		}{
+			{"Creating diff file", manager.CreateDiffFile},
+			{"Detecting dependencies", manager.DetectDependencyIntegrations},
+			{"Tracing package calls", manager.TracePackageCalls},
+			{"Scanning application", manager.ScanApplication},
+			{"Instrumenting application", manager.InstrumentApplication},
+			{"Resolving unit tests", manager.ResolveUnitTests},
+			{"Adding required modules", manager.AddRequiredModules},
+			{"Writing diff file", func() error {
+				comment.WriteAll()
+				// Pass a callback to WriteDiff to receive granular progress updates.
+				// This callback updates the UI with the name of the file currently being written,
+				// avoiding a "stalled" UI during this potentially long-running step.
+				return manager.WriteDiff(func(msg string) {
+					updates <- progressMsg{desc: msg}
+				})
+			}},
+		}
+
+		for _, step := range steps {
+			updates <- progressMsg{desc: step.desc}
+			if err := step.fn(); err != nil {
+				updates <- errMsg(err)
+				return
+			}
+		}
+
+		updates <- completedMsg{}
+		close(updates)
+	}()
+
+	initialM := initialModel(packagePath, outputFile)
+	initialM.sub = updates
+
+	finalModel, err := tea.NewProgram(initialM).Run()
+	if err != nil {
+		fmt.Printf("Error running program: %v\n", err)
+		os.Exit(1)
+	}
+
+	if m, ok := finalModel.(model); ok {
+		if m.err != nil {
+			os.Exit(1)
+		}
+		if m.done {
+			fmt.Printf("\nDone! Changes written to: %s\nTip: Apply these changes with: git apply %s\n", m.outputFile, m.outputFile)
+		}
+	}
+}
+
+func (m model) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
+	switch msg := msg.(type) {
+	case tea.KeyMsg:
+		if msg.Type == tea.KeyCtrlC {
+			return m, tea.Quit
+		}
+	case spinner.TickMsg:
+		var cmd tea.Cmd
+		m.spinner, cmd = m.spinner.Update(msg)
+		return m, cmd
+	case pkgLoadedMsg:
+		m.packages = msg
+		m.stepDesc = "Starting instrumentation..."
+		return m, waitForNext(m.sub)
+	case progressMsg:
+		m.currentStep++
+		m.stepDesc = msg.desc
+		// We just update the description, no progress bar to update
+		return m, waitForNext(m.sub)
+	case errMsg:
+		m.err = msg
+		return m, tea.Quit
+	case completedMsg:
+		m.done = true
+		return m, tea.Quit
+	}
+
+	// If we are strictly in Init, we should return the initial batch.
+	// But Update is called for every message.
+	// If the message is none of the above (shouldn't happen usually), we return nil.
+
+	// Wait, we need to ensure the first waitForNext is called.
+	// We can do it in a special "start" message or just include it in Init.
+	return m, nil
+}
+
+func waitForNext(sub chan tea.Msg) tea.Cmd {
+	if sub == nil {
+		return nil
+	}
+	return func() tea.Msg {
+		msg, ok := <-sub
+		if !ok {
+			return nil
+		}
+		return msg
+	}
+}
+
+func (m model) View() string {
+	if m.err != nil {
+		return fmt.Sprintf("\nError: %v\n", m.err)
+	}
+
+	pad := strings.Repeat(" ", padding(m.stepDesc, 30))
+
+	// Simple spinner view for all phases
+	return fmt.Sprintf("\n %s %s%s\n\n", m.spinner.View(), m.stepDesc, pad)
+}
+
+func padding(s string, width int) int {
+	l := len(s)
+	if l > width {
+		return 0
+	}
+	return width - l
 }
 
 func init() {
-	instrumentCmd.Flags().BoolVarP(&debug, "debug", "D", defaultDebug, "enable debugging output")
 	instrumentCmd.Flags().StringVarP(&diffFile, "output", "o", defaultOutputFilePath, "specify diff output file path")
+	instrumentCmd.Flags().StringVarP(&excludeDirs, "exclude", "e", "", "comma-separated list of folders to exclude from instrumentation")
 	cobra.MarkFlagFilename(instrumentCmd.Flags(), "output", ".diff") // for file completion
 
 	rootCmd.AddCommand(instrumentCmd)