root.go

  1package cmd
  2
  3import (
  4	"bytes"
  5	"context"
  6	"errors"
  7	"fmt"
  8	"io"
  9	"log/slog"
 10	"os"
 11	"path/filepath"
 12	"strconv"
 13	"strings"
 14
 15	tea "charm.land/bubbletea/v2"
 16	"charm.land/lipgloss/v2"
 17	"github.com/charmbracelet/colorprofile"
 18	"github.com/charmbracelet/crush/internal/app"
 19	"github.com/charmbracelet/crush/internal/config"
 20	"github.com/charmbracelet/crush/internal/db"
 21	"github.com/charmbracelet/crush/internal/event"
 22	"github.com/charmbracelet/crush/internal/projects"
 23	"github.com/charmbracelet/crush/internal/ui/common"
 24	ui "github.com/charmbracelet/crush/internal/ui/model"
 25	"github.com/charmbracelet/crush/internal/version"
 26	"github.com/charmbracelet/fang"
 27	uv "github.com/charmbracelet/ultraviolet"
 28	"github.com/charmbracelet/x/ansi"
 29	"github.com/charmbracelet/x/exp/charmtone"
 30	"github.com/charmbracelet/x/term"
 31	"github.com/spf13/cobra"
 32)
 33
 34func init() {
 35	rootCmd.PersistentFlags().StringP("cwd", "c", "", "Current working directory")
 36	rootCmd.PersistentFlags().StringP("data-dir", "D", "", "Custom crush data directory")
 37	rootCmd.PersistentFlags().BoolP("debug", "d", false, "Debug")
 38	rootCmd.Flags().BoolP("help", "h", false, "Help")
 39	rootCmd.Flags().BoolP("yolo", "y", false, "Automatically accept all permissions (dangerous mode)")
 40
 41	rootCmd.AddCommand(
 42		runCmd,
 43		dirsCmd,
 44		projectsCmd,
 45		updateProvidersCmd,
 46		logsCmd,
 47		schemaCmd,
 48		loginCmd,
 49		statsCmd,
 50	)
 51}
 52
 53var rootCmd = &cobra.Command{
 54	Use:   "crush",
 55	Short: "An AI assistant for software development",
 56	Long:  "An AI assistant for software development and similar tasks with direct access to the terminal",
 57	Example: `
 58# Run in interactive mode
 59crush
 60
 61# Run with debug logging
 62crush -d
 63
 64# Run with debug logging in a specific directory
 65crush -d -c /path/to/project
 66
 67# Run with custom data directory
 68crush -D /path/to/custom/.crush
 69
 70# Print version
 71crush -v
 72
 73# Run a single non-interactive prompt
 74crush run "Explain the use of context in Go"
 75
 76# Run in dangerous mode (auto-accept all permissions)
 77crush -y
 78  `,
 79	RunE: func(cmd *cobra.Command, args []string) error {
 80		app, err := setupAppWithProgressBar(cmd)
 81		if err != nil {
 82			return err
 83		}
 84		defer app.Shutdown()
 85
 86		event.AppInitialized()
 87
 88		// Set up the TUI.
 89		var env uv.Environ = os.Environ()
 90
 91		com := common.DefaultCommon(app)
 92		model := ui.New(com)
 93
 94		program := tea.NewProgram(
 95			model,
 96			tea.WithEnvironment(env),
 97			tea.WithContext(cmd.Context()),
 98			tea.WithFilter(ui.MouseEventFilter), // Filter mouse events based on focus state
 99		)
100		go app.Subscribe(program)
101
102		if _, err := program.Run(); err != nil {
103			event.Error(err)
104			slog.Error("TUI run error", "error", err)
105			return errors.New("Crush crashed. If metrics are enabled, we were notified about it. If you'd like to report it, please copy the stacktrace above and open an issue at https://github.com/charmbracelet/crush/issues/new?template=bug.yml") //nolint:staticcheck
106		}
107		return nil
108	},
109}
110
111var heartbit = lipgloss.NewStyle().Foreground(charmtone.Dolly).SetString(`
112    ▄▄▄▄▄▄▄▄    ▄▄▄▄▄▄▄▄
113  ███████████  ███████████
114████████████████████████████
115████████████████████████████
116██████████▀██████▀██████████
117██████████ ██████ ██████████
118▀▀██████▄████▄▄████▄██████▀▀
119  ████████████████████████
120    ████████████████████
121       ▀▀██████████▀▀
122           ▀▀▀▀▀▀
123`)
124
125// copied from cobra:
126const defaultVersionTemplate = `{{with .DisplayName}}{{printf "%s " .}}{{end}}{{printf "version %s" .Version}}
127`
128
129func Execute() {
130	// NOTE: very hacky: we create a colorprofile writer with STDOUT, then make
131	// it forward to a bytes.Buffer, write the colored heartbit to it, and then
132	// finally prepend it in the version template.
133	// Unfortunately cobra doesn't give us a way to set a function to handle
134	// printing the version, and PreRunE runs after the version is already
135	// handled, so that doesn't work either.
136	// This is the only way I could find that works relatively well.
137	if term.IsTerminal(os.Stdout.Fd()) {
138		var b bytes.Buffer
139		w := colorprofile.NewWriter(os.Stdout, os.Environ())
140		w.Forward = &b
141		_, _ = w.WriteString(heartbit.String())
142		rootCmd.SetVersionTemplate(b.String() + "\n" + defaultVersionTemplate)
143	}
144	if err := fang.Execute(
145		context.Background(),
146		rootCmd,
147		fang.WithVersion(version.Version),
148		fang.WithNotifySignal(os.Interrupt),
149	); err != nil {
150		os.Exit(1)
151	}
152}
153
154// supportsProgressBar tries to determine whether the current terminal supports
155// progress bars by looking into environment variables.
156func supportsProgressBar() bool {
157	if !term.IsTerminal(os.Stderr.Fd()) {
158		return false
159	}
160	termProg := os.Getenv("TERM_PROGRAM")
161	_, isWindowsTerminal := os.LookupEnv("WT_SESSION")
162
163	return isWindowsTerminal || strings.Contains(strings.ToLower(termProg), "ghostty")
164}
165
166func setupAppWithProgressBar(cmd *cobra.Command) (*app.App, error) {
167	app, err := setupApp(cmd)
168	if err != nil {
169		return nil, err
170	}
171
172	// Check if progress bar is enabled in config (defaults to true if nil)
173	progressEnabled := app.Config().Options.Progress == nil || *app.Config().Options.Progress
174	if progressEnabled && supportsProgressBar() {
175		_, _ = fmt.Fprintf(os.Stderr, ansi.SetIndeterminateProgressBar)
176		defer func() { _, _ = fmt.Fprintf(os.Stderr, ansi.ResetProgressBar) }()
177	}
178
179	return app, nil
180}
181
182// setupApp handles the common setup logic for both interactive and non-interactive modes.
183// It returns the app instance, config, cleanup function, and any error.
184func setupApp(cmd *cobra.Command) (*app.App, error) {
185	debug, _ := cmd.Flags().GetBool("debug")
186	yolo, _ := cmd.Flags().GetBool("yolo")
187	dataDir, _ := cmd.Flags().GetString("data-dir")
188	ctx := cmd.Context()
189
190	cwd, err := ResolveCwd(cmd)
191	if err != nil {
192		return nil, err
193	}
194
195	cfg, err := config.Init(cwd, dataDir, debug)
196	if err != nil {
197		return nil, err
198	}
199
200	if cfg.Permissions == nil {
201		cfg.Permissions = &config.Permissions{}
202	}
203	cfg.Permissions.SkipRequests = yolo
204
205	if err := createDotCrushDir(cfg.Options.DataDirectory); err != nil {
206		return nil, err
207	}
208
209	// Register this project in the centralized projects list.
210	if err := projects.Register(cwd, cfg.Options.DataDirectory); err != nil {
211		slog.Warn("Failed to register project", "error", err)
212		// Non-fatal: continue even if registration fails
213	}
214
215	// Connect to DB; this will also run migrations.
216	conn, err := db.Connect(ctx, cfg.Options.DataDirectory)
217	if err != nil {
218		return nil, err
219	}
220
221	appInstance, err := app.New(ctx, conn, cfg)
222	if err != nil {
223		slog.Error("Failed to create app instance", "error", err)
224		return nil, err
225	}
226
227	if shouldEnableMetrics(cfg) {
228		event.Init()
229	}
230
231	return appInstance, nil
232}
233
234func shouldEnableMetrics(cfg *config.Config) bool {
235	if v, _ := strconv.ParseBool(os.Getenv("CRUSH_DISABLE_METRICS")); v {
236		return false
237	}
238	if v, _ := strconv.ParseBool(os.Getenv("DO_NOT_TRACK")); v {
239		return false
240	}
241	if cfg.Options.DisableMetrics {
242		return false
243	}
244	return true
245}
246
247func MaybePrependStdin(prompt string) (string, error) {
248	if term.IsTerminal(os.Stdin.Fd()) {
249		return prompt, nil
250	}
251	fi, err := os.Stdin.Stat()
252	if err != nil {
253		return prompt, err
254	}
255	// Check if stdin is a named pipe ( | ) or regular file ( < ).
256	if fi.Mode()&os.ModeNamedPipe == 0 && !fi.Mode().IsRegular() {
257		return prompt, nil
258	}
259	bts, err := io.ReadAll(os.Stdin)
260	if err != nil {
261		return prompt, err
262	}
263	return string(bts) + "\n\n" + prompt, nil
264}
265
266func ResolveCwd(cmd *cobra.Command) (string, error) {
267	cwd, _ := cmd.Flags().GetString("cwd")
268	if cwd != "" {
269		err := os.Chdir(cwd)
270		if err != nil {
271			return "", fmt.Errorf("failed to change directory: %v", err)
272		}
273		return cwd, nil
274	}
275	cwd, err := os.Getwd()
276	if err != nil {
277		return "", fmt.Errorf("failed to get current working directory: %v", err)
278	}
279	return cwd, nil
280}
281
282func createDotCrushDir(dir string) error {
283	if err := os.MkdirAll(dir, 0o700); err != nil {
284		return fmt.Errorf("failed to create data directory: %q %w", dir, err)
285	}
286
287	gitIgnorePath := filepath.Join(dir, ".gitignore")
288	if _, err := os.Stat(gitIgnorePath); os.IsNotExist(err) {
289		if err := os.WriteFile(gitIgnorePath, []byte("*\n"), 0o644); err != nil {
290			return fmt.Errorf("failed to create .gitignore file: %q %w", gitIgnorePath, err)
291		}
292	}
293
294	return nil
295}