1package cmd
  2
  3import (
  4	"bytes"
  5	"context"
  6	"errors"
  7	"fmt"
  8	"io"
  9	"io/fs"
 10	"log/slog"
 11	"net/url"
 12	"os"
 13	"os/exec"
 14	"path/filepath"
 15	"regexp"
 16	"strconv"
 17	"strings"
 18	"time"
 19
 20	tea "charm.land/bubbletea/v2"
 21	fang "charm.land/fang/v2"
 22	"charm.land/lipgloss/v2"
 23	"github.com/charmbracelet/colorprofile"
 24	"github.com/charmbracelet/crush/internal/client"
 25	"github.com/charmbracelet/crush/internal/config"
 26	"github.com/charmbracelet/crush/internal/event"
 27	"github.com/charmbracelet/crush/internal/proto"
 28	"github.com/charmbracelet/crush/internal/server"
 29	"github.com/charmbracelet/crush/internal/ui/common"
 30	ui "github.com/charmbracelet/crush/internal/ui/model"
 31	"github.com/charmbracelet/crush/internal/version"
 32	"github.com/charmbracelet/crush/internal/workspace"
 33	uv "github.com/charmbracelet/ultraviolet"
 34	"github.com/charmbracelet/x/exp/charmtone"
 35	"github.com/charmbracelet/x/term"
 36	"github.com/spf13/cobra"
 37)
 38
 39var clientHost string
 40
 41func init() {
 42	rootCmd.PersistentFlags().StringP("cwd", "c", "", "Current working directory")
 43	rootCmd.PersistentFlags().StringP("data-dir", "D", "", "Custom crush data directory")
 44	rootCmd.PersistentFlags().BoolP("debug", "d", false, "Debug")
 45	rootCmd.PersistentFlags().BoolP("yolo", "y", false, "Automatically accept all permissions (dangerous mode)")
 46	rootCmd.PersistentFlags().StringVarP(&clientHost, "host", "H", server.DefaultHost(), "Connect to a specific crush server host (for advanced users)")
 47	rootCmd.Flags().BoolP("help", "h", false, "Help")
 48
 49	rootCmd.AddCommand(
 50		runCmd,
 51		dirsCmd,
 52		projectsCmd,
 53		updateProvidersCmd,
 54		logsCmd,
 55		schemaCmd,
 56		loginCmd,
 57		statsCmd,
 58		sessionCmd,
 59	)
 60}
 61
 62var rootCmd = &cobra.Command{
 63	Use:   "crush",
 64	Short: "A terminal-first AI assistant for software development",
 65	Long:  "A glamorous, terminal-first AI assistant for software development and adjacent tasks",
 66	Example: `
 67# Run in interactive mode
 68crush
 69
 70# Run non-interactively
 71crush run "Guess my 5 favorite Pokémon"
 72
 73# Run a non-interactively with pipes and redirection
 74cat README.md | crush run "make this more glamorous" > GLAMOROUS_README.md
 75
 76# Run with debug logging in a specific directory
 77crush --debug --cwd /path/to/project
 78
 79# Run in yolo mode (auto-accept all permissions; use with care)
 80crush --yolo
 81
 82# Run with custom data directory
 83crush --data-dir /path/to/custom/.crush
 84  `,
 85	RunE: func(cmd *cobra.Command, args []string) error {
 86		c, ws, cleanup, err := connectToServer(cmd)
 87		if err != nil {
 88			return err
 89		}
 90		defer cleanup()
 91
 92		event.AppInitialized()
 93
 94		clientWs := workspace.NewClientWorkspace(c, *ws)
 95
 96		if ws.Config.IsConfigured() {
 97			if err := clientWs.InitCoderAgent(cmd.Context()); err != nil {
 98				slog.Error("Failed to initialize coder agent", "error", err)
 99			}
100		}
101
102		com := common.DefaultCommon(clientWs)
103		model := ui.New(com)
104
105		var env uv.Environ = os.Environ()
106		program := tea.NewProgram(
107			model,
108			tea.WithEnvironment(env),
109			tea.WithContext(cmd.Context()),
110			tea.WithFilter(ui.MouseEventFilter),
111		)
112		go clientWs.Subscribe(program)
113
114		if _, err := program.Run(); err != nil {
115			event.Error(err)
116			slog.Error("TUI run error", "error", err)
117			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
118		}
119		return nil
120	},
121}
122
123var heartbit = lipgloss.NewStyle().Foreground(charmtone.Dolly).SetString(`
124    ▄▄▄▄▄▄▄▄    ▄▄▄▄▄▄▄▄
125  ███████████  ███████████
126████████████████████████████
127████████████████████████████
128██████████▀██████▀██████████
129██████████ ██████ ██████████
130▀▀██████▄████▄▄████▄██████▀▀
131  ████████████████████████
132    ████████████████████
133       ▀▀██████████▀▀
134           ▀▀▀▀▀▀
135`)
136
137// copied from cobra:
138const defaultVersionTemplate = `{{with .DisplayName}}{{printf "%s " .}}{{end}}{{printf "version %s" .Version}}
139`
140
141func Execute() {
142	// NOTE: very hacky: we create a colorprofile writer with STDOUT, then make
143	// it forward to a bytes.Buffer, write the colored heartbit to it, and then
144	// finally prepend it in the version template.
145	// Unfortunately cobra doesn't give us a way to set a function to handle
146	// printing the version, and PreRunE runs after the version is already
147	// handled, so that doesn't work either.
148	// This is the only way I could find that works relatively well.
149	if term.IsTerminal(os.Stdout.Fd()) {
150		var b bytes.Buffer
151		w := colorprofile.NewWriter(os.Stdout, os.Environ())
152		w.Forward = &b
153		_, _ = w.WriteString(heartbit.String())
154		rootCmd.SetVersionTemplate(b.String() + "\n" + defaultVersionTemplate)
155	}
156	if err := fang.Execute(
157		context.Background(),
158		rootCmd,
159		fang.WithVersion(version.Version),
160		fang.WithNotifySignal(os.Interrupt),
161	); err != nil {
162		os.Exit(1)
163	}
164}
165
166// supportsProgressBar tries to determine whether the current terminal supports
167// progress bars by looking into environment variables.
168func supportsProgressBar() bool {
169	if !term.IsTerminal(os.Stderr.Fd()) {
170		return false
171	}
172	termProg := os.Getenv("TERM_PROGRAM")
173	_, isWindowsTerminal := os.LookupEnv("WT_SESSION")
174
175	return isWindowsTerminal || strings.Contains(strings.ToLower(termProg), "ghostty")
176}
177
178// connectToServer ensures the server is running, creates a client and
179// workspace, and returns a cleanup function that deletes the workspace.
180func connectToServer(cmd *cobra.Command) (*client.Client, *proto.Workspace, func(), error) {
181	hostURL, err := server.ParseHostURL(clientHost)
182	if err != nil {
183		return nil, nil, nil, fmt.Errorf("invalid host URL: %v", err)
184	}
185
186	if err := ensureServer(cmd, hostURL); err != nil {
187		return nil, nil, nil, err
188	}
189
190	debug, _ := cmd.Flags().GetBool("debug")
191	yolo, _ := cmd.Flags().GetBool("yolo")
192	dataDir, _ := cmd.Flags().GetString("data-dir")
193	ctx := cmd.Context()
194
195	cwd, err := ResolveCwd(cmd)
196	if err != nil {
197		return nil, nil, nil, err
198	}
199
200	c, err := client.NewClient(cwd, hostURL.Scheme, hostURL.Host)
201	if err != nil {
202		return nil, nil, nil, err
203	}
204
205	wsReq := proto.Workspace{
206		Path:    cwd,
207		DataDir: dataDir,
208		Debug:   debug,
209		YOLO:    yolo,
210		Version: version.Version,
211		Env:     os.Environ(),
212	}
213
214	ws, err := c.CreateWorkspace(ctx, wsReq)
215	if err != nil {
216		// The server socket may exist before the HTTP handler is ready.
217		// Retry a few times with a short backoff.
218		for range 5 {
219			select {
220			case <-ctx.Done():
221				return nil, nil, nil, ctx.Err()
222			case <-time.After(200 * time.Millisecond):
223			}
224			ws, err = c.CreateWorkspace(ctx, wsReq)
225			if err == nil {
226				break
227			}
228		}
229		if err != nil {
230			return nil, nil, nil, fmt.Errorf("failed to create workspace: %v", err)
231		}
232	}
233
234	if shouldEnableMetrics(ws.Config) {
235		event.Init()
236	}
237
238	cleanup := func() { _ = c.DeleteWorkspace(context.Background(), ws.ID) }
239	return c, ws, cleanup, nil
240}
241
242// ensureServer auto-starts a detached server if the socket file does not
243// exist. When the socket exists, it verifies that the running server
244// version matches the client; on mismatch it shuts down the old server
245// and starts a fresh one.
246func ensureServer(cmd *cobra.Command, hostURL *url.URL) error {
247	switch hostURL.Scheme {
248	case "unix", "npipe":
249		needsStart := false
250		if _, err := os.Stat(hostURL.Host); err != nil && errors.Is(err, fs.ErrNotExist) {
251			needsStart = true
252		} else if err == nil {
253			if err := restartIfStale(cmd, hostURL); err != nil {
254				slog.Warn("Failed to check server version, restarting", "error", err)
255				needsStart = true
256			}
257		}
258
259		if needsStart {
260			if err := startDetachedServer(cmd); err != nil {
261				return err
262			}
263		}
264
265		var err error
266		for range 10 {
267			_, err = os.Stat(hostURL.Host)
268			if err == nil {
269				break
270			}
271			select {
272			case <-cmd.Context().Done():
273				return cmd.Context().Err()
274			case <-time.After(100 * time.Millisecond):
275			}
276		}
277		if err != nil {
278			return fmt.Errorf("failed to initialize crush server: %v", err)
279		}
280	}
281
282	return nil
283}
284
285// restartIfStale checks whether the running server matches the current
286// client version. When they differ, it sends a shutdown command and
287// removes the stale socket so the caller can start a fresh server.
288func restartIfStale(cmd *cobra.Command, hostURL *url.URL) error {
289	c, err := client.NewClient("", hostURL.Scheme, hostURL.Host)
290	if err != nil {
291		return err
292	}
293	vi, err := c.VersionInfo(cmd.Context())
294	if err != nil {
295		return err
296	}
297	if vi.Version == version.Version {
298		return nil
299	}
300	slog.Info("Server version mismatch, restarting",
301		"server", vi.Version,
302		"client", version.Version,
303	)
304	_ = c.ShutdownServer(cmd.Context())
305	// Give the old process a moment to release the socket.
306	for range 20 {
307		if _, err := os.Stat(hostURL.Host); errors.Is(err, fs.ErrNotExist) {
308			break
309		}
310		select {
311		case <-cmd.Context().Done():
312			return cmd.Context().Err()
313		case <-time.After(100 * time.Millisecond):
314		}
315	}
316	// Force-remove if the socket is still lingering.
317	_ = os.Remove(hostURL.Host)
318	return nil
319}
320
321var safeNameRegexp = regexp.MustCompile(`[^a-zA-Z0-9._-]`)
322
323func startDetachedServer(cmd *cobra.Command) error {
324	exe, err := os.Executable()
325	if err != nil {
326		return fmt.Errorf("failed to get executable path: %v", err)
327	}
328
329	safeClientHost := safeNameRegexp.ReplaceAllString(clientHost, "_")
330	chDir := filepath.Join(config.GlobalCacheDir(), "server-"+safeClientHost)
331	if err := os.MkdirAll(chDir, 0o700); err != nil {
332		return fmt.Errorf("failed to create server working directory: %v", err)
333	}
334
335	cmdArgs := []string{"server"}
336	if clientHost != server.DefaultHost() {
337		cmdArgs = append(cmdArgs, "--host", clientHost)
338	}
339
340	c := exec.CommandContext(cmd.Context(), exe, cmdArgs...)
341	stdoutPath := filepath.Join(chDir, "stdout.log")
342	stderrPath := filepath.Join(chDir, "stderr.log")
343	detachProcess(c)
344
345	stdout, err := os.Create(stdoutPath)
346	if err != nil {
347		return fmt.Errorf("failed to create stdout log file: %v", err)
348	}
349	defer stdout.Close()
350	c.Stdout = stdout
351
352	stderr, err := os.Create(stderrPath)
353	if err != nil {
354		return fmt.Errorf("failed to create stderr log file: %v", err)
355	}
356	defer stderr.Close()
357	c.Stderr = stderr
358
359	if err := c.Start(); err != nil {
360		return fmt.Errorf("failed to start crush server: %v", err)
361	}
362
363	if err := c.Process.Release(); err != nil {
364		return fmt.Errorf("failed to detach crush server process: %v", err)
365	}
366
367	return nil
368}
369
370func shouldEnableMetrics(cfg *config.Config) bool {
371	if v, _ := strconv.ParseBool(os.Getenv("CRUSH_DISABLE_METRICS")); v {
372		return false
373	}
374	if v, _ := strconv.ParseBool(os.Getenv("DO_NOT_TRACK")); v {
375		return false
376	}
377	if cfg.Options.DisableMetrics {
378		return false
379	}
380	return true
381}
382
383func MaybePrependStdin(prompt string) (string, error) {
384	if term.IsTerminal(os.Stdin.Fd()) {
385		return prompt, nil
386	}
387	fi, err := os.Stdin.Stat()
388	if err != nil {
389		return prompt, err
390	}
391	// Check if stdin is a named pipe ( | ) or regular file ( < ).
392	if fi.Mode()&os.ModeNamedPipe == 0 && !fi.Mode().IsRegular() {
393		return prompt, nil
394	}
395	bts, err := io.ReadAll(os.Stdin)
396	if err != nil {
397		return prompt, err
398	}
399	return string(bts) + "\n\n" + prompt, nil
400}
401
402func ResolveCwd(cmd *cobra.Command) (string, error) {
403	cwd, _ := cmd.Flags().GetString("cwd")
404	if cwd != "" {
405		err := os.Chdir(cwd)
406		if err != nil {
407			return "", fmt.Errorf("failed to change directory: %v", err)
408		}
409		return cwd, nil
410	}
411	cwd, err := os.Getwd()
412	if err != nil {
413		return "", fmt.Errorf("failed to get current working directory: %v", err)
414	}
415	return cwd, nil
416}