feat: stream content in non-interactive mode (#133)

Carlos Alexandro Becker created

Change summary

cmd/root.go               | 18 -------
internal/app/app.go       | 64 +++++++++++++++++-----------
internal/format/format.go | 91 ----------------------------------------
3 files changed, 40 insertions(+), 133 deletions(-)

Detailed changes

cmd/root.go 🔗

@@ -12,7 +12,6 @@ import (
 	"github.com/charmbracelet/crush/internal/app"
 	"github.com/charmbracelet/crush/internal/config"
 	"github.com/charmbracelet/crush/internal/db"
-	"github.com/charmbracelet/crush/internal/format"
 	"github.com/charmbracelet/crush/internal/llm/agent"
 	"github.com/charmbracelet/crush/internal/log"
 	"github.com/charmbracelet/crush/internal/tui"
@@ -52,14 +51,8 @@ to assist developers in writing, debugging, and understanding code directly from
 		debug, _ := cmd.Flags().GetBool("debug")
 		cwd, _ := cmd.Flags().GetString("cwd")
 		prompt, _ := cmd.Flags().GetString("prompt")
-		outputFormat, _ := cmd.Flags().GetString("output-format")
 		quiet, _ := cmd.Flags().GetBool("quiet")
 
-		// Validate format option
-		if !format.IsValid(outputFormat) {
-			return fmt.Errorf("invalid format option: %s\n%s", outputFormat, format.GetHelpText())
-		}
-
 		if cwd != "" {
 			err := os.Chdir(cwd)
 			if err != nil {
@@ -109,7 +102,7 @@ to assist developers in writing, debugging, and understanding code directly from
 		// Non-interactive mode
 		if prompt != "" {
 			// Run non-interactive flow using the App method
-			return app.RunNonInteractive(ctx, prompt, outputFormat, quiet)
+			return app.RunNonInteractive(ctx, prompt, quiet)
 		}
 
 		// Set up the TUI
@@ -164,17 +157,8 @@ func init() {
 	rootCmd.Flags().BoolP("debug", "d", false, "Debug")
 	rootCmd.Flags().StringP("prompt", "p", "", "Prompt to run in non-interactive mode")
 
-	// Add format flag with validation logic
-	rootCmd.Flags().StringP("output-format", "f", format.Text.String(),
-		"Output format for non-interactive mode (text, json)")
-
 	// Add quiet flag to hide spinner in non-interactive mode
 	rootCmd.Flags().BoolP("quiet", "q", false, "Hide spinner in non-interactive mode")
-
-	// Register custom validation for the format flag
-	rootCmd.RegisterFlagCompletionFunc("output-format", func(cmd *cobra.Command, args []string, toComplete string) ([]string, cobra.ShellCompDirective) {
-		return format.SupportedFormats, cobra.ShellCompDirectiveNoFileComp
-	})
 }
 
 func maybePrependStdin(prompt string) (string, error) {

internal/app/app.go 🔗

@@ -92,7 +92,7 @@ func New(ctx context.Context, conn *sql.DB, cfg *config.Config) (*App, error) {
 }
 
 // RunNonInteractive handles the execution flow when a prompt is provided via CLI flag.
-func (a *App) RunNonInteractive(ctx context.Context, prompt string, outputFormat string, quiet bool) error {
+func (a *App) RunNonInteractive(ctx context.Context, prompt string, quiet bool) error {
 	slog.Info("Running in non-interactive mode")
 
 	// Start spinner if not in quiet mode
@@ -100,8 +100,15 @@ func (a *App) RunNonInteractive(ctx context.Context, prompt string, outputFormat
 	if !quiet {
 		spinner = format.NewSpinner(ctx, "Generating")
 		spinner.Start()
-		defer spinner.Stop()
 	}
+	// Helper function to stop spinner once
+	stopSpinner := func() {
+		if !quiet && spinner != nil {
+			spinner.Stop()
+			spinner = nil
+		}
+	}
+	defer stopSpinner()
 
 	const maxPromptLengthForTitle = 100
 	titlePrefix := "Non-interactive: "
@@ -128,35 +135,42 @@ func (a *App) RunNonInteractive(ctx context.Context, prompt string, outputFormat
 		return fmt.Errorf("failed to start agent processing stream: %w", err)
 	}
 
-	result := <-done
+	messageEvents := a.Messages.Subscribe(ctx)
+	readBts := 0
 
-	// Stop spinner before printing output
-	if !quiet && spinner != nil {
-		spinner.Stop()
-	}
+	for {
+		select {
+		case result := <-done:
+			stopSpinner()
+
+			if result.Error != nil {
+				if errors.Is(result.Error, context.Canceled) || errors.Is(result.Error, agent.ErrRequestCancelled) {
+					slog.Info("Agent processing cancelled", "session_id", sess.ID)
+					return nil
+				}
+				return fmt.Errorf("agent processing failed: %w", result.Error)
+			}
+
+			part := result.Message.Content().String()[readBts:]
+			fmt.Println(part)
 
-	if result.Error != nil {
-		if errors.Is(result.Error, context.Canceled) || errors.Is(result.Error, agent.ErrRequestCancelled) {
-			slog.Info("Agent processing cancelled", "session_id", sess.ID)
+			slog.Info("Non-interactive run completed", "session_id", sess.ID)
 			return nil
-		}
-		return fmt.Errorf("agent processing failed: %w", result.Error)
-	}
 
-	// Get the text content from the response
-	content := "No content available"
-	if result.Message.Content().String() != "" {
-		content = result.Message.Content().String()
-	}
+		case event := <-messageEvents:
+			msg := event.Payload
+			if msg.SessionID == sess.ID && msg.Role == message.Assistant && len(msg.Parts) > 0 {
+				stopSpinner()
+				part := msg.Content().String()[readBts:]
+				fmt.Print(part)
+				readBts += len(part)
+			}
 
-	out, err := format.FormatOutput(content, outputFormat)
-	if err != nil {
-		return err
+		case <-ctx.Done():
+			stopSpinner()
+			return ctx.Err()
+		}
 	}
-
-	fmt.Println(out)
-	slog.Info("Non-interactive run completed", "session_id", sess.ID)
-	return nil
 }
 
 func (app *App) UpdateAgentModel() error {

internal/format/format.go 🔗

@@ -1,91 +0,0 @@
-package format
-
-import (
-	"encoding/json"
-	"fmt"
-	"strings"
-)
-
-// OutputFormat represents the output format type for non-interactive mode
-type OutputFormat string
-
-const (
-	// Text format outputs the AI response as plain text.
-	Text OutputFormat = "text"
-
-	// JSON format outputs the AI response wrapped in a JSON object.
-	JSON OutputFormat = "json"
-)
-
-// String returns the string representation of the OutputFormat
-func (f OutputFormat) String() string {
-	return string(f)
-}
-
-// SupportedFormats is a list of all supported output formats as strings
-var SupportedFormats = []string{
-	string(Text),
-	string(JSON),
-}
-
-// Parse converts a string to an OutputFormat
-func Parse(s string) (OutputFormat, error) {
-	s = strings.ToLower(strings.TrimSpace(s))
-
-	switch s {
-	case string(Text):
-		return Text, nil
-	case string(JSON):
-		return JSON, nil
-	default:
-		return "", fmt.Errorf("invalid format: %s", s)
-	}
-}
-
-// IsValid checks if the provided format string is supported
-func IsValid(s string) bool {
-	_, err := Parse(s)
-	return err == nil
-}
-
-// GetHelpText returns a formatted string describing all supported formats
-func GetHelpText() string {
-	return fmt.Sprintf(`Supported output formats:
-- %s: Plain text output (default)
-- %s: Output wrapped in a JSON object`,
-		Text, JSON)
-}
-
-// FormatOutput formats the AI response according to the specified format
-func FormatOutput(content string, formatStr string) (string, error) {
-	format, err := Parse(formatStr)
-	if err != nil {
-		format = Text
-	}
-
-	switch format {
-	case JSON:
-		return formatAsJSON(content)
-	case Text:
-		fallthrough
-	default:
-		return content, nil
-	}
-}
-
-// formatAsJSON wraps the content in a simple JSON object
-func formatAsJSON(content string) (string, error) {
-	// Use the JSON package to properly escape the content
-	response := struct {
-		Response string `json:"response"`
-	}{
-		Response: content,
-	}
-
-	jsonBytes, err := json.MarshalIndent(response, "", "  ")
-	if err != nil {
-		return "", fmt.Errorf("failed to marshal output into JSON: %w", err)
-	}
-
-	return string(jsonBytes), nil
-}