run.go

  1package cmd
  2
  3import (
  4	"context"
  5	"fmt"
  6	"log/slog"
  7	"os"
  8	"os/signal"
  9	"strings"
 10	"time"
 11
 12	"charm.land/lipgloss/v2"
 13	"charm.land/log/v2"
 14	"github.com/charmbracelet/crush/internal/client"
 15	"github.com/charmbracelet/crush/internal/config"
 16	"github.com/charmbracelet/crush/internal/event"
 17	"github.com/charmbracelet/crush/internal/format"
 18	"github.com/charmbracelet/crush/internal/proto"
 19	"github.com/charmbracelet/crush/internal/pubsub"
 20	"github.com/charmbracelet/crush/internal/session"
 21	"github.com/charmbracelet/crush/internal/ui/anim"
 22	"github.com/charmbracelet/crush/internal/ui/styles"
 23	"github.com/charmbracelet/x/ansi"
 24	"github.com/charmbracelet/x/exp/charmtone"
 25	"github.com/charmbracelet/x/term"
 26	"github.com/spf13/cobra"
 27)
 28
 29var runCmd = &cobra.Command{
 30	Use:   "run [prompt...]",
 31	Short: "Run a single non-interactive prompt",
 32	Long: `Run a single prompt in non-interactive mode and exit.
 33The prompt can be provided as arguments or piped from stdin.`,
 34	Example: `
 35# Run a simple prompt
 36crush run "Guess my 5 favorite Pokรฉmon"
 37
 38# Pipe input from stdin
 39curl https://charm.land | crush run "Summarize this website"
 40
 41# Read from a file
 42crush run "What is this code doing?" <<< prrr.go
 43
 44# Redirect output to a file
 45crush run "Generate a hot README for this project" > MY_HOT_README.md
 46
 47# Run in quiet mode (hide the spinner)
 48crush run --quiet "Generate a README for this project"
 49
 50# Run in verbose mode (show logs)
 51crush run --verbose "Generate a README for this project"
 52
 53# Continue a previous session
 54crush run --session {session-id} "Follow up on your last response"
 55
 56# Continue the most recent session
 57crush run --continue "Follow up on your last response"
 58
 59  `,
 60	RunE: func(cmd *cobra.Command, args []string) error {
 61		var (
 62			quiet, _      = cmd.Flags().GetBool("quiet")
 63			verbose, _    = cmd.Flags().GetBool("verbose")
 64			largeModel, _ = cmd.Flags().GetString("model")
 65			smallModel, _ = cmd.Flags().GetString("small-model")
 66			sessionID, _  = cmd.Flags().GetString("session")
 67			useLast, _    = cmd.Flags().GetBool("continue")
 68		)
 69
 70		// Cancel on SIGINT or SIGTERM.
 71		ctx, cancel := signal.NotifyContext(context.Background(), os.Interrupt, os.Kill)
 72		defer cancel()
 73
 74		c, ws, cleanup, err := connectToServer(cmd)
 75		if err != nil {
 76			return err
 77		}
 78		defer cleanup()
 79
 80		if sessionID != "" {
 81			sess, err := resolveSessionByID(ctx, c, ws.ID, sessionID)
 82			if err != nil {
 83				return err
 84			}
 85			sessionID = sess.ID
 86		}
 87
 88		if !ws.Config.IsConfigured() {
 89			return fmt.Errorf("no providers configured - please run 'crush' to set up a provider interactively")
 90		}
 91
 92		if verbose {
 93			slog.SetDefault(slog.New(log.New(os.Stderr)))
 94		}
 95
 96		prompt := strings.Join(args, " ")
 97
 98		prompt, err = MaybePrependStdin(prompt)
 99		if err != nil {
100			slog.Error("Failed to read from stdin", "error", err)
101			return err
102		}
103
104		if prompt == "" {
105			return fmt.Errorf("no prompt provided")
106		}
107
108		event.SetNonInteractive(true)
109		event.AppInitialized()
110
111		switch {
112		case sessionID != "":
113			event.SetContinueBySessionID(true)
114		case useLast:
115			event.SetContinueLastSession(true)
116		}
117
118		return runNonInteractive(ctx, c, ws, prompt, largeModel, smallModel, quiet || verbose, sessionID, useLast)
119	},
120}
121
122func init() {
123	runCmd.Flags().BoolP("quiet", "q", false, "Hide spinner")
124	runCmd.Flags().BoolP("verbose", "v", false, "Show logs")
125	runCmd.Flags().StringP("model", "m", "", "Model to use. Accepts 'model' or 'provider/model' to disambiguate models with the same name across providers")
126	runCmd.Flags().String("small-model", "", "Small model to use. If not provided, uses the default small model for the provider")
127	runCmd.Flags().StringP("session", "s", "", "Continue a previous session by ID")
128	runCmd.Flags().BoolP("continue", "C", false, "Continue the most recent session")
129	runCmd.MarkFlagsMutuallyExclusive("session", "continue")
130}
131
132// runNonInteractive executes the agent via the server and streams output
133// to stdout.
134func runNonInteractive(
135	ctx context.Context,
136	c *client.Client,
137	ws *proto.Workspace,
138	prompt, largeModel, smallModel string,
139	hideSpinner bool,
140	continueSessionID string,
141	useLast bool,
142) error {
143	slog.Info("Running in non-interactive mode")
144
145	ctx, cancel := context.WithCancel(ctx)
146	defer cancel()
147
148	if largeModel != "" || smallModel != "" {
149		if err := overrideModels(ctx, c, ws, largeModel, smallModel); err != nil {
150			return fmt.Errorf("failed to override models: %w", err)
151		}
152	}
153
154	var (
155		spinner   *format.Spinner
156		stdoutTTY bool
157		stderrTTY bool
158		stdinTTY  bool
159		progress  bool
160	)
161
162	stdoutTTY = term.IsTerminal(os.Stdout.Fd())
163	stderrTTY = term.IsTerminal(os.Stderr.Fd())
164	stdinTTY = term.IsTerminal(os.Stdin.Fd())
165	progress = ws.Config.Options.Progress == nil || *ws.Config.Options.Progress
166
167	if !hideSpinner && stderrTTY {
168		t := styles.DefaultStyles()
169
170		hasDarkBG := true
171		if stdinTTY && stdoutTTY {
172			hasDarkBG = lipgloss.HasDarkBackground(os.Stdin, os.Stdout)
173		}
174		defaultFG := lipgloss.LightDark(hasDarkBG)(charmtone.Pepper, t.FgBase)
175
176		spinner = format.NewSpinner(ctx, cancel, anim.Settings{
177			Size:        10,
178			Label:       "Generating",
179			LabelColor:  defaultFG,
180			GradColorA:  t.Primary,
181			GradColorB:  t.Secondary,
182			CycleColors: true,
183		})
184		spinner.Start()
185	}
186
187	stopSpinner := func() {
188		if !hideSpinner && spinner != nil {
189			spinner.Stop()
190			spinner = nil
191		}
192	}
193
194	// Wait for the agent to become ready (MCP init, etc).
195	if err := waitForAgent(ctx, c, ws.ID); err != nil {
196		stopSpinner()
197		return fmt.Errorf("agent not ready: %w", err)
198	}
199
200	// Force-update agent models so MCP tools are loaded.
201	if err := c.UpdateAgent(ctx, ws.ID); err != nil {
202		slog.Warn("Failed to update agent", "error", err)
203	}
204
205	defer stopSpinner()
206
207	sess, err := resolveSession(ctx, c, ws.ID, continueSessionID, useLast)
208	if err != nil {
209		return fmt.Errorf("failed to resolve session: %w", err)
210	}
211	if continueSessionID != "" || useLast {
212		slog.Info("Continuing session for non-interactive run", "session_id", sess.ID)
213	} else {
214		slog.Info("Created session for non-interactive run", "session_id", sess.ID)
215	}
216
217	events, err := c.SubscribeEvents(ctx, ws.ID)
218	if err != nil {
219		return fmt.Errorf("failed to subscribe to events: %w", err)
220	}
221
222	if err := c.SendMessage(ctx, ws.ID, sess.ID, prompt); err != nil {
223		return fmt.Errorf("failed to send message: %w", err)
224	}
225
226	messageReadBytes := make(map[string]int)
227	var printed bool
228
229	defer func() {
230		if progress && stderrTTY {
231			_, _ = fmt.Fprintf(os.Stderr, ansi.ResetProgressBar)
232		}
233		_, _ = fmt.Fprintln(os.Stdout)
234	}()
235
236	for {
237		if progress && stderrTTY {
238			_, _ = fmt.Fprintf(os.Stderr, ansi.SetIndeterminateProgressBar)
239		}
240
241		select {
242		case ev, ok := <-events:
243			if !ok {
244				stopSpinner()
245				return nil
246			}
247
248			switch e := ev.(type) {
249			case pubsub.Event[proto.Message]:
250				msg := e.Payload
251				if msg.SessionID != sess.ID || msg.Role != proto.Assistant || len(msg.Parts) == 0 {
252					continue
253				}
254				stopSpinner()
255
256				content := msg.Content().String()
257				readBytes := messageReadBytes[msg.ID]
258
259				if len(content) < readBytes {
260					slog.Error("Non-interactive: message content shorter than read bytes",
261						"message_length", len(content), "read_bytes", readBytes)
262					return fmt.Errorf("message content is shorter than read bytes: %d < %d", len(content), readBytes)
263				}
264
265				part := content[readBytes:]
266				if readBytes == 0 {
267					part = strings.TrimLeft(part, " \t")
268				}
269				if printed || strings.TrimSpace(part) != "" {
270					printed = true
271					fmt.Fprint(os.Stdout, part)
272				}
273				messageReadBytes[msg.ID] = len(content)
274
275				if msg.IsFinished() {
276					return nil
277				}
278
279			case pubsub.Event[proto.AgentEvent]:
280				if e.Payload.Error != nil {
281					stopSpinner()
282					return fmt.Errorf("agent error: %w", e.Payload.Error)
283				}
284			}
285
286		case <-ctx.Done():
287			stopSpinner()
288			return ctx.Err()
289		}
290	}
291}
292
293// waitForAgent polls GetAgentInfo until the agent is ready, with a
294// timeout.
295func waitForAgent(ctx context.Context, c *client.Client, wsID string) error {
296	timeout := time.After(30 * time.Second)
297	for {
298		info, err := c.GetAgentInfo(ctx, wsID)
299		if err == nil && info.IsReady {
300			return nil
301		}
302		select {
303		case <-timeout:
304			if err != nil {
305				return fmt.Errorf("timeout waiting for agent: %w", err)
306			}
307			return fmt.Errorf("timeout waiting for agent readiness")
308		case <-ctx.Done():
309			return ctx.Err()
310		case <-time.After(200 * time.Millisecond):
311		}
312	}
313}
314
315// overrideModels resolves model strings and updates the workspace
316// configuration via the server.
317func overrideModels(
318	ctx context.Context,
319	c *client.Client,
320	ws *proto.Workspace,
321	largeModel, smallModel string,
322) error {
323	cfg, err := c.GetConfig(ctx, ws.ID)
324	if err != nil {
325		return fmt.Errorf("failed to get config: %w", err)
326	}
327
328	providers := cfg.Providers.Copy()
329
330	largeMatches, smallMatches := findModelMatches(providers, largeModel, smallModel)
331
332	var largeProviderID string
333
334	if largeModel != "" {
335		found, err := validateModelMatches(largeMatches, largeModel, "large")
336		if err != nil {
337			return err
338		}
339		largeProviderID = found.provider
340		slog.Info("Overriding large model", "provider", found.provider, "model", found.modelID)
341		if err := c.UpdatePreferredModel(ctx, ws.ID, config.ScopeWorkspace, config.SelectedModelTypeLarge, config.SelectedModel{
342			Provider: found.provider,
343			Model:    found.modelID,
344		}); err != nil {
345			return fmt.Errorf("failed to set large model: %w", err)
346		}
347	}
348
349	switch {
350	case smallModel != "":
351		found, err := validateModelMatches(smallMatches, smallModel, "small")
352		if err != nil {
353			return err
354		}
355		slog.Info("Overriding small model", "provider", found.provider, "model", found.modelID)
356		if err := c.UpdatePreferredModel(ctx, ws.ID, config.ScopeWorkspace, config.SelectedModelTypeSmall, config.SelectedModel{
357			Provider: found.provider,
358			Model:    found.modelID,
359		}); err != nil {
360			return fmt.Errorf("failed to set small model: %w", err)
361		}
362
363	case largeModel != "":
364		sm, err := c.GetDefaultSmallModel(ctx, ws.ID, largeProviderID)
365		if err != nil {
366			slog.Warn("Failed to get default small model", "error", err)
367		} else if sm != nil {
368			if err := c.UpdatePreferredModel(ctx, ws.ID, config.ScopeWorkspace, config.SelectedModelTypeSmall, *sm); err != nil {
369				return fmt.Errorf("failed to set small model: %w", err)
370			}
371		}
372	}
373
374	return c.UpdateAgent(ctx, ws.ID)
375}
376
377type modelMatch struct {
378	provider string
379	modelID  string
380}
381
382// findModelMatches searches providers for matching large/small model
383// strings.
384func findModelMatches(providers map[string]config.ProviderConfig, largeModel, smallModel string) ([]modelMatch, []modelMatch) {
385	largeFilter, largeID := parseModelString(largeModel)
386	smallFilter, smallID := parseModelString(smallModel)
387
388	var largeMatches, smallMatches []modelMatch
389	for name, provider := range providers {
390		if provider.Disable {
391			continue
392		}
393		for _, m := range provider.Models {
394			if matchesModel(largeID, largeFilter, m.ID, name) {
395				largeMatches = append(largeMatches, modelMatch{provider: name, modelID: m.ID})
396			}
397			if matchesModel(smallID, smallFilter, m.ID, name) {
398				smallMatches = append(smallMatches, modelMatch{provider: name, modelID: m.ID})
399			}
400		}
401	}
402	return largeMatches, smallMatches
403}
404
405// parseModelString splits "provider/model" into (provider, model) or
406// ("", model).
407func parseModelString(s string) (string, string) {
408	if s == "" {
409		return "", ""
410	}
411	if idx := strings.Index(s, "/"); idx >= 0 {
412		return s[:idx], s[idx+1:]
413	}
414	return "", s
415}
416
417// matchesModel returns true if the model ID matches the filter
418// criteria.
419func matchesModel(wantID, wantProvider, modelID, providerName string) bool {
420	if wantID == "" {
421		return false
422	}
423	if wantProvider != "" && wantProvider != providerName {
424		return false
425	}
426	return strings.EqualFold(modelID, wantID)
427}
428
429// validateModelMatches ensures exactly one match exists.
430func validateModelMatches(matches []modelMatch, modelID, label string) (modelMatch, error) {
431	switch {
432	case len(matches) == 0:
433		return modelMatch{}, fmt.Errorf("%s model %q not found", label, modelID)
434	case len(matches) > 1:
435		names := make([]string, len(matches))
436		for i, m := range matches {
437			names[i] = m.provider
438		}
439		return modelMatch{}, fmt.Errorf(
440			"%s model: model %q found in multiple providers: %s. Please specify provider using 'provider/model' format",
441			label, modelID, strings.Join(names, ", "),
442		)
443	}
444	return matches[0], nil
445}
446
447// resolveSession returns the session to use for a non-interactive run.
448// If continueSessionID is set it fetches that session; if useLast is set it
449// returns the most recently updated top-level session; otherwise it creates a
450// new one.
451func resolveSession(ctx context.Context, c *client.Client, wsID, continueSessionID string, useLast bool) (*session.Session, error) {
452	switch {
453	case continueSessionID != "":
454		sess, err := c.GetSession(ctx, wsID, continueSessionID)
455		if err != nil {
456			return nil, fmt.Errorf("session not found: %s", continueSessionID)
457		}
458		if sess.ParentSessionID != "" {
459			return nil, fmt.Errorf("cannot continue a child session: %s", continueSessionID)
460		}
461		return sess, nil
462
463	case useLast:
464		sessions, err := c.ListSessions(ctx, wsID)
465		if err != nil || len(sessions) == 0 {
466			return nil, fmt.Errorf("no sessions found to continue")
467		}
468		last := sessions[0]
469		for _, s := range sessions[1:] {
470			if s.UpdatedAt > last.UpdatedAt && s.ParentSessionID == "" {
471				last = s
472			}
473		}
474		return &last, nil
475
476	default:
477		return c.CreateSession(ctx, wsID, "non-interactive")
478	}
479}
480
481// resolveSessionByID resolves a session ID that may be a full UUID or a hash
482// prefix returned by crush session list.
483func resolveSessionByID(ctx context.Context, c *client.Client, wsID, id string) (*session.Session, error) {
484	if sess, err := c.GetSession(ctx, wsID, id); err == nil {
485		return sess, nil
486	}
487
488	sessions, err := c.ListSessions(ctx, wsID)
489	if err != nil {
490		return nil, err
491	}
492
493	var matches []session.Session
494	for _, s := range sessions {
495		hash := session.HashID(s.ID)
496		if hash == id || strings.HasPrefix(hash, id) {
497			matches = append(matches, s)
498		}
499	}
500
501	switch len(matches) {
502	case 0:
503		return nil, fmt.Errorf("session %q not found", id)
504	case 1:
505		return &matches[0], nil
506	default:
507		return nil, fmt.Errorf("session ID %q is ambiguous (%d matches)", id, len(matches))
508	}
509}