root.go

  1package cmd
  2
  3import (
  4	"context"
  5	"fmt"
  6	"io"
  7	"log/slog"
  8	"os"
  9	"time"
 10
 11	tea "github.com/charmbracelet/bubbletea/v2"
 12	"github.com/charmbracelet/crush/internal/app"
 13	"github.com/charmbracelet/crush/internal/config"
 14	"github.com/charmbracelet/crush/internal/db"
 15	"github.com/charmbracelet/crush/internal/format"
 16	"github.com/charmbracelet/crush/internal/llm/agent"
 17	"github.com/charmbracelet/crush/internal/log"
 18	"github.com/charmbracelet/crush/internal/tui"
 19	"github.com/charmbracelet/crush/internal/version"
 20	"github.com/charmbracelet/fang"
 21	"github.com/charmbracelet/x/term"
 22	"github.com/spf13/cobra"
 23)
 24
 25var rootCmd = &cobra.Command{
 26	Use:   "crush",
 27	Short: "Terminal-based AI assistant for software development",
 28	Long: `Crush is a powerful terminal-based AI assistant that helps with software development tasks.
 29It provides an interactive chat interface with AI capabilities, code analysis, and LSP integration
 30to assist developers in writing, debugging, and understanding code directly from the terminal.`,
 31	Example: `
 32  # Run in interactive mode
 33  crush
 34
 35  # Run with debug logging
 36  crush -d
 37
 38  # Run with debug slog.in a specific directory
 39  crush -d -c /path/to/project
 40
 41  # Print version
 42  crush -v
 43
 44  # Run a single non-interactive prompt
 45  crush -p "Explain the use of context in Go"
 46
 47  # Run a single non-interactive prompt with JSON output format
 48  crush -p "Explain the use of context in Go" -f json
 49  `,
 50	RunE: func(cmd *cobra.Command, args []string) error {
 51		// Load the config
 52		debug, _ := cmd.Flags().GetBool("debug")
 53		cwd, _ := cmd.Flags().GetString("cwd")
 54		prompt, _ := cmd.Flags().GetString("prompt")
 55		outputFormat, _ := cmd.Flags().GetString("output-format")
 56		quiet, _ := cmd.Flags().GetBool("quiet")
 57
 58		// Validate format option
 59		if !format.IsValid(outputFormat) {
 60			return fmt.Errorf("invalid format option: %s\n%s", outputFormat, format.GetHelpText())
 61		}
 62
 63		if cwd != "" {
 64			err := os.Chdir(cwd)
 65			if err != nil {
 66				return fmt.Errorf("failed to change directory: %v", err)
 67			}
 68		}
 69		if cwd == "" {
 70			c, err := os.Getwd()
 71			if err != nil {
 72				return fmt.Errorf("failed to get current working directory: %v", err)
 73			}
 74			cwd = c
 75		}
 76
 77		cfg, err := config.Init(cwd, debug)
 78		if err != nil {
 79			return err
 80		}
 81
 82		// Create main context for the application
 83		ctx, cancel := context.WithCancel(context.Background())
 84		defer cancel()
 85
 86		// Connect DB, this will also run migrations
 87		conn, err := db.Connect(ctx, cfg.Options.DataDirectory)
 88		if err != nil {
 89			return err
 90		}
 91
 92		app, err := app.New(ctx, conn, cfg)
 93		if err != nil {
 94			slog.Error(fmt.Sprintf("Failed to create app instance: %v", err))
 95			return err
 96		}
 97		// Defer shutdown here so it runs for both interactive and non-interactive modes
 98		defer app.Shutdown()
 99
100		// Initialize MCP tools early for both modes
101		initMCPTools(ctx, app, cfg)
102
103		prompt, err = maybePrependStdin(prompt)
104		if err != nil {
105			slog.Error(fmt.Sprintf("Failed to read from stdin: %v", err))
106			return err
107		}
108
109		// Non-interactive mode
110		if prompt != "" {
111			// Run non-interactive flow using the App method
112			return app.RunNonInteractive(ctx, prompt, outputFormat, quiet)
113		}
114
115		// Set up the TUI
116		program := tea.NewProgram(
117			tui.New(app),
118			tea.WithAltScreen(),
119			tea.WithKeyReleases(),
120			tea.WithUniformKeyLayout(),
121		)
122
123		go app.Subscribe(program)
124
125		if _, err := program.Run(); err != nil {
126			slog.Error(fmt.Sprintf("TUI run error: %v", err))
127			return fmt.Errorf("TUI error: %v", err)
128		}
129		app.Shutdown()
130		return nil
131	},
132}
133
134func initMCPTools(ctx context.Context, app *app.App, cfg *config.Config) {
135	go func() {
136		defer log.RecoverPanic("MCP-goroutine", nil)
137
138		// Create a context with timeout for the initial MCP tools fetch
139		ctxWithTimeout, cancel := context.WithTimeout(ctx, 30*time.Second)
140		defer cancel()
141
142		// Set this up once with proper error handling
143		agent.GetMcpTools(ctxWithTimeout, app.Permissions, cfg)
144		slog.Info("MCP message handling goroutine exiting")
145	}()
146}
147
148func Execute() {
149	if err := fang.Execute(
150		context.Background(),
151		rootCmd,
152		fang.WithVersion(version.Version),
153	); err != nil {
154		os.Exit(1)
155	}
156}
157
158func init() {
159	rootCmd.PersistentFlags().StringP("cwd", "c", "", "Current working directory")
160
161	rootCmd.Flags().BoolP("help", "h", false, "Help")
162	rootCmd.Flags().BoolP("debug", "d", false, "Debug")
163	rootCmd.Flags().StringP("prompt", "p", "", "Prompt to run in non-interactive mode")
164
165	// Add format flag with validation logic
166	rootCmd.Flags().StringP("output-format", "f", format.Text.String(),
167		"Output format for non-interactive mode (text, json)")
168
169	// Add quiet flag to hide spinner in non-interactive mode
170	rootCmd.Flags().BoolP("quiet", "q", false, "Hide spinner in non-interactive mode")
171
172	// Register custom validation for the format flag
173	rootCmd.RegisterFlagCompletionFunc("output-format", func(cmd *cobra.Command, args []string, toComplete string) ([]string, cobra.ShellCompDirective) {
174		return format.SupportedFormats, cobra.ShellCompDirectiveNoFileComp
175	})
176}
177
178func maybePrependStdin(prompt string) (string, error) {
179	if term.IsTerminal(os.Stdin.Fd()) {
180		return prompt, nil
181	}
182	fi, err := os.Stdin.Stat()
183	if err != nil {
184		return prompt, err
185	}
186	if fi.Mode()&os.ModeNamedPipe == 0 {
187		return prompt, nil
188	}
189	bts, err := io.ReadAll(os.Stdin)
190	if err != nil {
191		return prompt, err
192	}
193	return string(bts) + "\n\n" + prompt, nil
194}