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/stringext"
 23	termutil "github.com/charmbracelet/crush/internal/term"
 24	"github.com/charmbracelet/crush/internal/tui"
 25	"github.com/charmbracelet/crush/internal/ui/common"
 26	ui "github.com/charmbracelet/crush/internal/ui/model"
 27	"github.com/charmbracelet/crush/internal/version"
 28	"github.com/charmbracelet/fang"
 29	uv "github.com/charmbracelet/ultraviolet"
 30	"github.com/charmbracelet/x/ansi"
 31	"github.com/charmbracelet/x/exp/charmtone"
 32	"github.com/charmbracelet/x/term"
 33	"github.com/spf13/cobra"
 34)
 35
 36func init() {
 37	rootCmd.PersistentFlags().StringP("cwd", "c", "", "Current working directory")
 38	rootCmd.PersistentFlags().StringP("data-dir", "D", "", "Custom crush data directory")
 39	rootCmd.PersistentFlags().BoolP("debug", "d", false, "Debug")
 40	rootCmd.Flags().BoolP("help", "h", false, "Help")
 41	rootCmd.Flags().BoolP("yolo", "y", false, "Automatically accept all permissions (dangerous mode)")
 42
 43	rootCmd.AddCommand(
 44		runCmd,
 45		dirsCmd,
 46		updateProvidersCmd,
 47		logsCmd,
 48		schemaCmd,
 49		loginCmd,
 50	)
 51}
 52
 53var rootCmd = &cobra.Command{
 54	Use:   "crush",
 55	Short: "Terminal-based AI assistant for software development",
 56	Long: `Crush is a powerful terminal-based AI assistant that helps with software development tasks.
 57It provides an interactive chat interface with AI capabilities, code analysis, and LSP integration
 58to assist developers in writing, debugging, and understanding code directly from the terminal.`,
 59	Example: `
 60# Run in interactive mode
 61crush
 62
 63# Run with debug logging
 64crush -d
 65
 66# Run with debug logging in a specific directory
 67crush -d -c /path/to/project
 68
 69# Run with custom data directory
 70crush -D /path/to/custom/.crush
 71
 72# Print version
 73crush -v
 74
 75# Run a single non-interactive prompt
 76crush run "Explain the use of context in Go"
 77
 78# Run in dangerous mode (auto-accept all permissions)
 79crush -y
 80  `,
 81	RunE: func(cmd *cobra.Command, args []string) error {
 82		app, err := setupAppWithProgressBar(cmd)
 83		if err != nil {
 84			return err
 85		}
 86		defer app.Shutdown()
 87
 88		event.AppInitialized()
 89
 90		// Set up the TUI.
 91		var env uv.Environ = os.Environ()
 92		com := common.DefaultCommon(app)
 93		ui := ui.New(com)
 94		ui.QueryVersion = shouldQueryTerminalVersion(env)
 95		program := tea.NewProgram(
 96			ui,
 97			tea.WithEnvironment(env),
 98			tea.WithContext(cmd.Context()),
 99			tea.WithFilter(tui.MouseEventFilter)) // Filter mouse events based on focus state
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	PostRun: func(cmd *cobra.Command, args []string) {
110		event.AppExited()
111	},
112}
113
114var heartbit = lipgloss.NewStyle().Foreground(charmtone.Dolly).SetString(`
115    ▄▄▄▄▄▄▄▄    ▄▄▄▄▄▄▄▄
116  ███████████  ███████████
117████████████████████████████
118████████████████████████████
119██████████▀██████▀██████████
120██████████ ██████ ██████████
121▀▀██████▄████▄▄████▄██████▀▀
122  ████████████████████████
123    ████████████████████
124       ▀▀██████████▀▀
125           ▀▀▀▀▀▀
126`)
127
128// copied from cobra:
129const defaultVersionTemplate = `{{with .DisplayName}}{{printf "%s " .}}{{end}}{{printf "version %s" .Version}}
130`
131
132func Execute() {
133	// NOTE: very hacky: we create a colorprofile writer with STDOUT, then make
134	// it forward to a bytes.Buffer, write the colored heartbit to it, and then
135	// finally prepend it in the version template.
136	// Unfortunately cobra doesn't give us a way to set a function to handle
137	// printing the version, and PreRunE runs after the version is already
138	// handled, so that doesn't work either.
139	// This is the only way I could find that works relatively well.
140	if term.IsTerminal(os.Stdout.Fd()) {
141		var b bytes.Buffer
142		w := colorprofile.NewWriter(os.Stdout, os.Environ())
143		w.Forward = &b
144		_, _ = w.WriteString(heartbit.String())
145		rootCmd.SetVersionTemplate(b.String() + "\n" + defaultVersionTemplate)
146	}
147	if err := fang.Execute(
148		context.Background(),
149		rootCmd,
150		fang.WithVersion(version.Version),
151		fang.WithNotifySignal(os.Interrupt),
152	); err != nil {
153		os.Exit(1)
154	}
155}
156
157func setupAppWithProgressBar(cmd *cobra.Command) (*app.App, error) {
158	if termutil.SupportsProgressBar() {
159		_, _ = fmt.Fprintf(os.Stderr, ansi.SetIndeterminateProgressBar)
160		defer func() { _, _ = fmt.Fprintf(os.Stderr, ansi.ResetProgressBar) }()
161	}
162
163	return setupApp(cmd)
164}
165
166// setupApp handles the common setup logic for both interactive and non-interactive modes.
167// It returns the app instance, config, cleanup function, and any error.
168func setupApp(cmd *cobra.Command) (*app.App, error) {
169	debug, _ := cmd.Flags().GetBool("debug")
170	yolo, _ := cmd.Flags().GetBool("yolo")
171	dataDir, _ := cmd.Flags().GetString("data-dir")
172	ctx := cmd.Context()
173
174	cwd, err := ResolveCwd(cmd)
175	if err != nil {
176		return nil, err
177	}
178
179	cfg, err := config.Init(cwd, dataDir, debug)
180	if err != nil {
181		return nil, err
182	}
183
184	if cfg.Permissions == nil {
185		cfg.Permissions = &config.Permissions{}
186	}
187	cfg.Permissions.SkipRequests = yolo
188
189	if err := createDotCrushDir(cfg.Options.DataDirectory); err != nil {
190		return nil, err
191	}
192
193	// Connect to DB; this will also run migrations.
194	conn, err := db.Connect(ctx, cfg.Options.DataDirectory)
195	if err != nil {
196		return nil, err
197	}
198
199	appInstance, err := app.New(ctx, conn, cfg)
200	if err != nil {
201		slog.Error("Failed to create app instance", "error", err)
202		return nil, err
203	}
204
205	if shouldEnableMetrics() {
206		event.Init()
207	}
208
209	return appInstance, nil
210}
211
212func shouldEnableMetrics() bool {
213	if v, _ := strconv.ParseBool(os.Getenv("CRUSH_DISABLE_METRICS")); v {
214		return false
215	}
216	if v, _ := strconv.ParseBool(os.Getenv("DO_NOT_TRACK")); v {
217		return false
218	}
219	if config.Get().Options.DisableMetrics {
220		return false
221	}
222	return true
223}
224
225func MaybePrependStdin(prompt string) (string, error) {
226	if term.IsTerminal(os.Stdin.Fd()) {
227		return prompt, nil
228	}
229	fi, err := os.Stdin.Stat()
230	if err != nil {
231		return prompt, err
232	}
233	if fi.Mode()&os.ModeNamedPipe == 0 {
234		return prompt, nil
235	}
236	bts, err := io.ReadAll(os.Stdin)
237	if err != nil {
238		return prompt, err
239	}
240	return string(bts) + "\n\n" + prompt, nil
241}
242
243func ResolveCwd(cmd *cobra.Command) (string, error) {
244	cwd, _ := cmd.Flags().GetString("cwd")
245	if cwd != "" {
246		err := os.Chdir(cwd)
247		if err != nil {
248			return "", fmt.Errorf("failed to change directory: %v", err)
249		}
250		return cwd, nil
251	}
252	cwd, err := os.Getwd()
253	if err != nil {
254		return "", fmt.Errorf("failed to get current working directory: %v", err)
255	}
256	return cwd, nil
257}
258
259func createDotCrushDir(dir string) error {
260	if err := os.MkdirAll(dir, 0o700); err != nil {
261		return fmt.Errorf("failed to create data directory: %q %w", dir, err)
262	}
263
264	gitIgnorePath := filepath.Join(dir, ".gitignore")
265	if _, err := os.Stat(gitIgnorePath); os.IsNotExist(err) {
266		if err := os.WriteFile(gitIgnorePath, []byte("*\n"), 0o644); err != nil {
267			return fmt.Errorf("failed to create .gitignore file: %q %w", gitIgnorePath, err)
268		}
269	}
270
271	return nil
272}
273
274func shouldQueryTerminalVersion(env uv.Environ) bool {
275	termType := env.Getenv("TERM")
276	termProg, okTermProg := env.LookupEnv("TERM_PROGRAM")
277	_, okSSHTTY := env.LookupEnv("SSH_TTY")
278	return (!okTermProg && !okSSHTTY) ||
279		(!strings.Contains(termProg, "Apple") && !okSSHTTY) ||
280		// Terminals that do support XTVERSION.
281		stringext.ContainsAny(termType, "alacritty", "ghostty", "kitty", "rio", "wezterm")
282}