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