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