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	"time"
 15
 16	tea "charm.land/bubbletea/v2"
 17	"charm.land/lipgloss/v2"
 18	"github.com/charmbracelet/colorprofile"
 19	"github.com/charmbracelet/crush/internal/app"
 20	"github.com/charmbracelet/crush/internal/config"
 21	"github.com/charmbracelet/crush/internal/db"
 22	"github.com/charmbracelet/crush/internal/event"
 23	"github.com/charmbracelet/crush/internal/stringext"
 24	termutil "github.com/charmbracelet/crush/internal/term"
 25	"github.com/charmbracelet/crush/internal/tui"
 26	tuiutil "github.com/charmbracelet/crush/internal/tui/util"
 27	"github.com/charmbracelet/crush/internal/update"
 28	"github.com/charmbracelet/crush/internal/version"
 29	"github.com/charmbracelet/fang"
 30	uv "github.com/charmbracelet/ultraviolet"
 31	"github.com/charmbracelet/x/ansi"
 32	"github.com/charmbracelet/x/exp/charmtone"
 33	"github.com/charmbracelet/x/term"
 34	"github.com/spf13/cobra"
 35)
 36
 37func init() {
 38	rootCmd.PersistentFlags().StringP("cwd", "c", "", "Current working directory")
 39	rootCmd.PersistentFlags().StringP("data-dir", "D", "", "Custom crush data directory")
 40	rootCmd.PersistentFlags().BoolP("debug", "d", false, "Debug")
 41	rootCmd.Flags().BoolP("help", "h", false, "Help")
 42	rootCmd.Flags().BoolP("yolo", "y", false, "Automatically accept all permissions (dangerous mode)")
 43
 44	rootCmd.AddCommand(
 45		runCmd,
 46		dirsCmd,
 47		updateCmd,
 48		updateProvidersCmd,
 49		logsCmd,
 50		schemaCmd,
 51		loginCmd,
 52	)
 53}
 54
 55var rootCmd = &cobra.Command{
 56	Use:   "crush",
 57	Short: "Terminal-based AI assistant for software development",
 58	Long: `Crush is a powerful terminal-based AI assistant that helps with software development tasks.
 59It provides an interactive chat interface with AI capabilities, code analysis, and LSP integration
 60to assist developers in writing, debugging, and understanding code directly from the terminal.`,
 61	Example: `
 62# Run in interactive mode
 63crush
 64
 65# Run with debug logging
 66crush -d
 67
 68# Run with debug logging in a specific directory
 69crush -d -c /path/to/project
 70
 71# Run with custom data directory
 72crush -D /path/to/custom/.crush
 73
 74# Print version
 75crush -v
 76
 77# Run a single non-interactive prompt
 78crush run "Explain the use of context in Go"
 79
 80# Run in dangerous mode (auto-accept all permissions)
 81crush -y
 82  `,
 83	RunE: func(cmd *cobra.Command, args []string) error {
 84		app, err := setupAppWithProgressBar(cmd)
 85		if err != nil {
 86			return err
 87		}
 88		defer app.Shutdown()
 89
 90		event.AppInitialized()
 91
 92		// Set up the TUI.
 93		var env uv.Environ = os.Environ()
 94		ui := tui.New(app)
 95		ui.QueryVersion = shouldQueryTerminalVersion(env)
 96
 97		program := tea.NewProgram(
 98			ui,
 99			tea.WithEnvironment(env),
100			tea.WithContext(cmd.Context()),
101			tea.WithFilter(tui.MouseEventFilter)) // Filter mouse events based on focus state
102		go app.Subscribe(program)
103
104		// Create a cancellable context for the update check that gets cancelled
105		// when the TUI exits.
106		updateCtx, cancelUpdate := context.WithCancel(cmd.Context())
107		defer cancelUpdate()
108
109		// Start async update check unless disabled.
110		go checkForUpdateAsync(updateCtx, program)
111
112		if _, err := program.Run(); err != nil {
113			event.Error(err)
114			slog.Error("TUI run error", "error", err)
115			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
116		}
117		return nil
118	},
119	PostRun: func(cmd *cobra.Command, args []string) {
120		event.AppExited()
121	},
122}
123
124var heartbit = lipgloss.NewStyle().Foreground(charmtone.Dolly).SetString(`
125    ▄▄▄▄▄▄▄▄    ▄▄▄▄▄▄▄▄
126  ███████████  ███████████
127████████████████████████████
128████████████████████████████
129██████████▀██████▀██████████
130██████████ ██████ ██████████
131▀▀██████▄████▄▄████▄██████▀▀
132  ████████████████████████
133    ████████████████████
134       ▀▀██████████▀▀
135           ▀▀▀▀▀▀
136`)
137
138// copied from cobra:
139const defaultVersionTemplate = `{{with .DisplayName}}{{printf "%s " .}}{{end}}{{printf "version %s" .Version}}
140`
141
142func Execute() {
143	// NOTE: very hacky: we create a colorprofile writer with STDOUT, then make
144	// it forward to a bytes.Buffer, write the colored heartbit to it, and then
145	// finally prepend it in the version template.
146	// Unfortunately cobra doesn't give us a way to set a function to handle
147	// printing the version, and PreRunE runs after the version is already
148	// handled, so that doesn't work either.
149	// This is the only way I could find that works relatively well.
150	versionTemplate := defaultVersionTemplate
151	if term.IsTerminal(os.Stdout.Fd()) {
152		var b bytes.Buffer
153		w := colorprofile.NewWriter(os.Stdout, os.Environ())
154		w.Forward = &b
155		_, _ = w.WriteString(heartbit.String())
156		versionTemplate = b.String() + "\n" + defaultVersionTemplate
157	}
158
159	// Check if version flag is present and add update notification if available.
160	if hasVersionFlag() {
161		if updateMsg := checkForUpdateSync(); updateMsg != "" {
162			versionTemplate += updateMsg
163		}
164	}
165
166	rootCmd.SetVersionTemplate(versionTemplate)
167
168	if err := fang.Execute(
169		context.Background(),
170		rootCmd,
171		fang.WithVersion(version.Version),
172		fang.WithNotifySignal(os.Interrupt),
173	); err != nil {
174		os.Exit(1)
175	}
176}
177
178// hasVersionFlag checks if the version flag is present in os.Args.
179func hasVersionFlag() bool {
180	for _, arg := range os.Args {
181		if arg == "-v" || arg == "--version" {
182			return true
183		}
184	}
185	return false
186}
187
188// isAutoUpdateDisabled checks if update checks are disabled via env var.
189// Config is not loaded at this point (called before Execute), so only env var is checked.
190func isAutoUpdateDisabled() bool {
191	if str, ok := os.LookupEnv("CRUSH_DISABLE_AUTO_UPDATE"); ok {
192		v, _ := strconv.ParseBool(str)
193		return v
194	}
195	return false
196}
197
198// checkForUpdateSync performs a synchronous update check with a short timeout.
199// Returns a formatted update message if an update is available, empty string otherwise.
200func checkForUpdateSync() string {
201	if isAutoUpdateDisabled() {
202		return ""
203	}
204
205	ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
206	defer cancel()
207
208	info, err := update.Check(ctx, version.Version, update.Default)
209	if err != nil || !info.Available() {
210		return ""
211	}
212
213	if info.IsDevelopment() {
214		return info.DevelopmentVersionBrief()
215	}
216
217	return fmt.Sprintf("\nUpdate available: v%s → v%s\nRun 'crush update apply' to install.\n", info.Current, info.Latest)
218}
219
220// checkForUpdateAsync checks for updates in the background and applies them if possible.
221func checkForUpdateAsync(ctx context.Context, program *tea.Program) {
222	// Check config (if loaded) or env var.
223	if isAutoUpdateDisabled() {
224		return
225	}
226	if cfg := config.Get(); cfg != nil && cfg.Options.DisableAutoUpdate {
227		return
228	}
229
230	checkCtx, cancel := context.WithTimeout(ctx, 2*time.Minute)
231	defer cancel()
232
233	info, err := update.Check(checkCtx, version.Version, update.Default)
234	if err != nil || !info.Available() || info.IsDevelopment() {
235		return
236	}
237
238	// Check if context was cancelled while checking.
239	if ctx.Err() != nil {
240		return
241	}
242
243	// Check install method.
244	method := update.DetectInstallMethod()
245	if !method.CanSelfUpdate() {
246		// Package manager install - show instructions.
247		program.Send(tuiutil.InfoMsg{
248			Type: tuiutil.InfoTypeUpdate,
249			Msg:  fmt.Sprintf("Update available: v%s → v%s. Run: %s", info.Current, info.Latest, method.UpdateInstructions()),
250			TTL:  30 * time.Second,
251		})
252		return
253	}
254
255	// Attempt self-update.
256	asset, err := update.FindAsset(info.Release.Assets)
257	if err != nil {
258		program.Send(tuiutil.InfoMsg{
259			Type: tuiutil.InfoTypeWarn,
260			Msg:  "Update available but failed to find asset. Run 'crush update' for details.",
261			TTL:  15 * time.Second,
262		})
263		return
264	}
265
266	// Check if context was cancelled before download.
267	if ctx.Err() != nil {
268		return
269	}
270
271	binaryPath, err := update.Download(checkCtx, asset, info.Release)
272	if err != nil {
273		// Don't show error message if context was cancelled (user exited).
274		if ctx.Err() != nil {
275			return
276		}
277		program.Send(tuiutil.InfoMsg{
278			Type: tuiutil.InfoTypeWarn,
279			Msg:  "Update download failed. Run 'crush update' for details.",
280			TTL:  15 * time.Second,
281		})
282		return
283	}
284	defer os.Remove(binaryPath)
285
286	// Check if context was cancelled before apply.
287	if ctx.Err() != nil {
288		return
289	}
290
291	if err := update.Apply(binaryPath); err != nil {
292		program.Send(tuiutil.InfoMsg{
293			Type: tuiutil.InfoTypeWarn,
294			Msg:  "Update failed to install. Run 'crush update' for details.",
295			TTL:  15 * time.Second,
296		})
297		return
298	}
299
300	// Success!
301	program.Send(tuiutil.InfoMsg{
302		Type: tuiutil.InfoTypeUpdate,
303		Msg:  fmt.Sprintf("Updated to v%s! Restart Crush to use the new version.", info.Latest),
304		TTL:  30 * time.Second,
305	})
306}
307
308func setupAppWithProgressBar(cmd *cobra.Command) (*app.App, error) {
309	if termutil.SupportsProgressBar() {
310		_, _ = fmt.Fprintf(os.Stderr, ansi.SetIndeterminateProgressBar)
311		defer func() { _, _ = fmt.Fprintf(os.Stderr, ansi.ResetProgressBar) }()
312	}
313
314	return setupApp(cmd)
315}
316
317// setupApp handles the common setup logic for both interactive and non-interactive modes.
318// It returns the app instance, config, cleanup function, and any error.
319func setupApp(cmd *cobra.Command) (*app.App, error) {
320	debug, _ := cmd.Flags().GetBool("debug")
321	yolo, _ := cmd.Flags().GetBool("yolo")
322	dataDir, _ := cmd.Flags().GetString("data-dir")
323	ctx := cmd.Context()
324
325	cwd, err := ResolveCwd(cmd)
326	if err != nil {
327		return nil, err
328	}
329
330	cfg, err := config.Init(cwd, dataDir, debug)
331	if err != nil {
332		return nil, err
333	}
334
335	if cfg.Permissions == nil {
336		cfg.Permissions = &config.Permissions{}
337	}
338	cfg.Permissions.SkipRequests = yolo
339
340	if err := createDotCrushDir(cfg.Options.DataDirectory); err != nil {
341		return nil, err
342	}
343
344	// Connect to DB; this will also run migrations.
345	conn, err := db.Connect(ctx, cfg.Options.DataDirectory)
346	if err != nil {
347		return nil, err
348	}
349
350	appInstance, err := app.New(ctx, conn, cfg)
351	if err != nil {
352		slog.Error("Failed to create app instance", "error", err)
353		return nil, err
354	}
355
356	if shouldEnableMetrics() {
357		event.Init()
358	}
359
360	return appInstance, nil
361}
362
363func shouldEnableMetrics() bool {
364	if v, _ := strconv.ParseBool(os.Getenv("CRUSH_DISABLE_METRICS")); v {
365		return false
366	}
367	if v, _ := strconv.ParseBool(os.Getenv("DO_NOT_TRACK")); v {
368		return false
369	}
370	if config.Get().Options.DisableMetrics {
371		return false
372	}
373	return true
374}
375
376func MaybePrependStdin(prompt string) (string, error) {
377	if term.IsTerminal(os.Stdin.Fd()) {
378		return prompt, nil
379	}
380	fi, err := os.Stdin.Stat()
381	if err != nil {
382		return prompt, err
383	}
384	if fi.Mode()&os.ModeNamedPipe == 0 {
385		return prompt, nil
386	}
387	bts, err := io.ReadAll(os.Stdin)
388	if err != nil {
389		return prompt, err
390	}
391	return string(bts) + "\n\n" + prompt, nil
392}
393
394func ResolveCwd(cmd *cobra.Command) (string, error) {
395	cwd, _ := cmd.Flags().GetString("cwd")
396	if cwd != "" {
397		err := os.Chdir(cwd)
398		if err != nil {
399			return "", fmt.Errorf("failed to change directory: %v", err)
400		}
401		return cwd, nil
402	}
403	cwd, err := os.Getwd()
404	if err != nil {
405		return "", fmt.Errorf("failed to get current working directory: %v", err)
406	}
407	return cwd, nil
408}
409
410func createDotCrushDir(dir string) error {
411	if err := os.MkdirAll(dir, 0o700); err != nil {
412		return fmt.Errorf("failed to create data directory: %q %w", dir, err)
413	}
414
415	gitIgnorePath := filepath.Join(dir, ".gitignore")
416	if _, err := os.Stat(gitIgnorePath); os.IsNotExist(err) {
417		if err := os.WriteFile(gitIgnorePath, []byte("*\n"), 0o644); err != nil {
418			return fmt.Errorf("failed to create .gitignore file: %q %w", gitIgnorePath, err)
419		}
420	}
421
422	return nil
423}
424
425func shouldQueryTerminalVersion(env uv.Environ) bool {
426	termType := env.Getenv("TERM")
427	termProg, okTermProg := env.LookupEnv("TERM_PROGRAM")
428	_, okSSHTTY := env.LookupEnv("SSH_TTY")
429	return (!okTermProg && !okSSHTTY) ||
430		(!strings.Contains(termProg, "Apple") && !okSSHTTY) ||
431		// Terminals that do support XTVERSION.
432		stringext.ContainsAny(termType, "alacritty", "ghostty", "kitty", "rio", "wezterm")
433}