Improve light mode colors and fix goal mode color bug

butterfish/butterfish.go

@@ -57,7 +57,8 @@ type ButterfishConfig struct {
 
 	// A list of context-specific styles drawn from the colorscheme
 	// These are what should actually be used during rendering
-	Styles *styles
+	Styles    *styles
+	ColorDark bool
 
 	// Path of yaml file from which to load LLM prompts
 	// Defaults to ~/.config/butterfish/prompts.yaml
@@ -70,7 +71,6 @@ type ButterfishConfig struct {
 	// Shell mode configuration
 	ShellMode               bool
 	ShellPluginMode         bool
-	ShellColorDark          bool
 	ShellBinary             string // path to the shell binary to use, e.g. /bin/zsh
 	ShellPromptModel        string // used when the user enters an explicit prompt
 	ShellLeavePromptAlone   bool   // don't try to edit the shell prompt

butterfish/commands.go

@@ -622,7 +622,11 @@ func (this *ButterfishCtx) Prompt(cmd *promptCommand) (*util.CompletionResponse,
 		termWidth, _, _ := term.GetSize(int(os.Stdout.Fd()))
 
 		if termWidth > 0 {
-			writer = util.NewStyleCodeblocksWriter(this.Out, termWidth, color, highlight)
+			colorScheme := "monokai"
+			if !this.Config.ColorDark {
+				colorScheme = "monokailight"
+			}
+			writer = util.NewStyleCodeblocksWriter(this.Out, termWidth, color, highlight, colorScheme)
 		}
 	} else if cmd.NoBackticks {
 		// this is an else because the code blocks writer will strip out backticks

butterfish/shell.go

@@ -60,8 +60,8 @@ var DarkShellColorScheme = &ShellColorScheme{
 	PromptGoalUnsafe: "\x1b[38;5;9m",
 	Command:          "\x1b[0m",
 	Autosuggest:      "\x1b[38;5;241m",
-	Answer:           "\x1b[38;5;221m",
-	AnswerHighlight:  "\x1b[38;5;204m",
+	Answer:           "\x1b[38;5;221m", // yellow
+	AnswerHighlight:  "\x1b[38;5;204m", // orange
 	GoalMode:         "\x1b[38;5;51m",
 	Error:            "\x1b[38;5;196m",
 }
@@ -72,8 +72,8 @@ var LightShellColorScheme = &ShellColorScheme{
 	PromptGoalUnsafe: "\x1b[38;5;9m",
 	Command:          "\x1b[0m",
 	Autosuggest:      "\x1b[38;5;241m",
-	Answer:           "\x1b[38;5;221m",
-	AnswerHighlight:  "\x1b[38;5;204m",
+	Answer:           "\x1b[38;5;18m", // Dark blue
+	AnswerHighlight:  "\x1b[38;5;6m",
 	GoalMode:         "\x1b[38;5;18m",
 	Error:            "\x1b[38;5;196m",
 }
@@ -366,29 +366,30 @@ type ShellState struct {
 	AutosuggestMaxTokens int
 
 	// The current state of the shell
-	State                int
-	GoalMode             bool
-	GoalModeBuffer       string
-	GoalModeGoal         string
-	GoalModeUnsafe       bool
-	ActiveFunction       string
-	PromptSuffixCounter  int
-	ChildOutReader       chan *byteMsg
-	ParentInReader       chan *byteMsg
-	CursorPosChan        chan *cursorPosition
-	PromptOutputChan     chan *util.CompletionResponse
-	PrintErrorChan       chan error
-	AutosuggestChan      chan *AutosuggestResult
-	History              *ShellHistory
-	PromptAnswerWriter   io.Writer
-	StyleWriter          *util.StyleCodeblocksWriter
-	Prompt               *ShellBuffer
-	PromptResponseCancel context.CancelFunc
-	Command              *ShellBuffer
-	TerminalWidth        int
-	Color                *ShellColorScheme
-	LastTabPassthrough   time.Time
-	parentInBuffer       []byte
+	State                  int
+	GoalMode               bool
+	GoalModeBuffer         string
+	GoalModeGoal           string
+	GoalModeUnsafe         bool
+	ActiveFunction         string
+	PromptSuffixCounter    int
+	ChildOutReader         chan *byteMsg
+	ParentInReader         chan *byteMsg
+	CursorPosChan          chan *cursorPosition
+	PromptOutputChan       chan *util.CompletionResponse
+	PrintErrorChan         chan error
+	AutosuggestChan        chan *AutosuggestResult
+	History                *ShellHistory
+	PromptAnswerWriter     io.Writer
+	PromptGoalAnswerWriter io.Writer
+	StyleWriter            *util.StyleCodeblocksWriter
+	Prompt                 *ShellBuffer
+	PromptResponseCancel   context.CancelFunc
+	Command                *ShellBuffer
+	TerminalWidth          int
+	Color                  *ShellColorScheme
+	LastTabPassthrough     time.Time
+	parentInBuffer         []byte
 	// these are used to estimate number of tokens
 	AutosuggestEncoder *tiktoken.Tiktoken
 	PromptEncoder      *tiktoken.Tiktoken
@@ -586,7 +587,7 @@ func (this *ButterfishCtx) ShellMultiplexer(
 	this.SetPS1(childIn)
 
 	colorScheme := DarkShellColorScheme
-	if !this.Config.ShellColorDark {
+	if !this.Config.ColorDark {
 		colorScheme = LightShellColorScheme
 	}
 
@@ -604,8 +605,22 @@ func (this *ButterfishCtx) ShellMultiplexer(
 	}
 
 	carriageReturnWriter := util.NewReplaceWriter(parentOut, "\n", "\r\n")
-	styleCodeblocksWriter := util.NewStyleCodeblocksWriter(carriageReturnWriter,
-		termWidth, colorScheme.Answer, colorScheme.AnswerHighlight)
+	codeblocksColorScheme := "monokai"
+	if !this.Config.ColorDark {
+		codeblocksColorScheme = "monokailight"
+	}
+	styleCodeblocksWriter := util.NewStyleCodeblocksWriter(
+		carriageReturnWriter,
+		termWidth,
+		colorScheme.Answer,
+		colorScheme.AnswerHighlight,
+		codeblocksColorScheme)
+	styleCodeblocksWriterGoal := util.NewStyleCodeblocksWriter(
+		carriageReturnWriter,
+		termWidth,
+		colorScheme.GoalMode,
+		colorScheme.AnswerHighlight,
+		codeblocksColorScheme)
 
 	sigwinch := make(chan os.Signal, 1)
 	signal.Notify(sigwinch, syscall.SIGWINCH)
@@ -618,28 +633,29 @@ func (this *ButterfishCtx) ShellMultiplexer(
 		this.Config.ShellMaxPromptTokens)
 
 	shellState := &ShellState{
-		Butterfish:           this,
-		ParentOut:            parentOut,
-		ChildIn:              childIn,
-		Sigwinch:             sigwinch,
-		State:                stateNormal,
-		ChildOutReader:       childOutReader,
-		ParentInReader:       parentInReader,
-		CursorPosChan:        parentPositionChan,
-		PrintErrorChan:       make(chan error, 8),
-		History:              NewShellHistory(),
-		PromptOutputChan:     make(chan *util.CompletionResponse),
-		PromptAnswerWriter:   styleCodeblocksWriter,
-		StyleWriter:          styleCodeblocksWriter,
-		Command:              NewShellBuffer(),
-		Prompt:               NewShellBuffer(),
-		TerminalWidth:        termWidth,
-		AutosuggestEnabled:   this.Config.ShellAutosuggestEnabled,
-		AutosuggestChan:      make(chan *AutosuggestResult),
-		Color:                colorScheme,
-		parentInBuffer:       []byte{},
-		PromptMaxTokens:      promptMaxTokens,
-		AutosuggestMaxTokens: autoSuggestMaxTokens,
+		Butterfish:             this,
+		ParentOut:              parentOut,
+		ChildIn:                childIn,
+		Sigwinch:               sigwinch,
+		State:                  stateNormal,
+		ChildOutReader:         childOutReader,
+		ParentInReader:         parentInReader,
+		CursorPosChan:          parentPositionChan,
+		PrintErrorChan:         make(chan error, 8),
+		History:                NewShellHistory(),
+		PromptOutputChan:       make(chan *util.CompletionResponse),
+		PromptAnswerWriter:     styleCodeblocksWriter,
+		PromptGoalAnswerWriter: styleCodeblocksWriterGoal,
+		StyleWriter:            styleCodeblocksWriter,
+		Command:                NewShellBuffer(),
+		Prompt:                 NewShellBuffer(),
+		TerminalWidth:          termWidth,
+		AutosuggestEnabled:     this.Config.ShellAutosuggestEnabled,
+		AutosuggestChan:        make(chan *AutosuggestResult),
+		Color:                  colorScheme,
+		parentInBuffer:         []byte{},
+		PromptMaxTokens:        promptMaxTokens,
+		AutosuggestMaxTokens:   autoSuggestMaxTokens,
 	}
 
 	shellState.Prompt.SetTerminalWidth(termWidth)
@@ -1000,7 +1016,7 @@ func (this *ShellState) ParentInput(ctx context.Context, data []byte) []byte {
 		if data[0] == 0x03 {
 			if this.GoalMode {
 				// Ctrl-C while in goal mode
-				fmt.Fprintf(this.PromptAnswerWriter, "\n%sExited goal mode.%s\n", this.Color.Answer, this.Color.Command)
+				fmt.Fprintf(this.PromptGoalAnswerWriter, "\n%sExited goal mode.%s\n", this.Color.Answer, this.Color.Command)
 				this.GoalMode = false
 			}
 
@@ -1279,7 +1295,7 @@ func (this *ShellState) GoalModeStart() {
 	}
 
 	this.GoalMode = true
-	fmt.Fprintf(this.PromptAnswerWriter, "%sGoal mode starting...%s\n", this.Color.Answer, this.Color.Command)
+	fmt.Fprintf(this.PromptGoalAnswerWriter, "%sGoal mode starting...%s\n", this.Color.Answer, this.Color.Command)
 	this.GoalModeGoal = goal
 	this.Prompt.Clear()
 
@@ -1359,7 +1375,7 @@ func (this *ShellState) GoalModeFunction(output *util.CompletionResponse) {
 			result = "FAILURE"
 		}
 
-		fmt.Fprintf(this.PromptAnswerWriter, "%sExited goal mode with %s.%s\n", this.Color.Answer, result, this.Color.Command)
+		fmt.Fprintf(this.PromptGoalAnswerWriter, "%sExited goal mode with %s.%s\n", this.Color.Answer, result, this.Color.Command)
 		this.GoalMode = false
 
 	case "":
@@ -1476,7 +1492,7 @@ func (this *ShellState) goalModePrompt(lastPrompt string) {
 	// we run this in a goroutine so that we can still receive input
 	// like Ctrl-C while waiting for the response
 	go CompletionRoutine(request, this.Butterfish.LLMClient,
-		this.PromptAnswerWriter, this.PromptOutputChan,
+		this.PromptGoalAnswerWriter, this.PromptOutputChan,
 		this.Color.GoalMode, this.Color.Error, this.StyleWriter)
 }
 

cmd/butterfish/main.go

@@ -77,6 +77,7 @@ type CliConfig struct {
 	Version      kong.VersionFlag `short:"V" help:"Print version information and exit."`
 	BaseURL      string           `short:"u" default:"https://api.openai.com/v1" help:"Base URL for OpenAI-compatible API. Enables local models with a compatible interface."`
 	TokenTimeout int              `short:"z" default:"10000" help:"Timeout before first prompt token is received and between individual tokens. In milliseconds."`
+	LightColor   bool             `short:"l" default:"false" help:"Light color mode, appropriate for a terminal with a white(ish) background"`
 
 	Shell struct {
 		Bin                       string `short:"b" help:"Shell to use (e.g. /bin/zsh), defaults to $SHELL."`
@@ -86,7 +87,6 @@ type CliConfig struct {
 		AutosuggestTimeout        int    `short:"t" default:"500" help:"Delay after typing before autosuggest (lower values trigger more calls and are more expensive). In milliseconds."`
 		NewlineAutosuggestTimeout int    `short:"T" default:"3500" help:"Timeout for autosuggest on a fresh line, i.e. before a command has started. Negative values disable. In milliseconds."`
 		NoCommandPrompt           bool   `short:"p" default:"false" help:"Don't change command prompt (shell PS1 variable). If not set, an emoji will be added to the prompt as a reminder you're in Shell Mode."`
-		LightColor                bool   `short:"l" default:"false" help:"Light color mode, appropriate for a terminal with a white(ish) background"`
 		MaxPromptTokens           int    `short:"P" default:"16384" help:"Maximum number of tokens, we restrict calls to this size regardless of model capabilities."`
 		MaxHistoryBlockTokens     int    `short:"H" default:"1024" help:"Maximum number of tokens of each block of history. For example, if a command has a very long output, it will be truncated to this length when sending the shell's history."`
 		MaxResponseTokens         int    `short:"R" default:"2048" help:"Maximum number of tokens in a response when prompting."`
@@ -232,7 +232,7 @@ func main() {
 		config.ShellAutosuggestModel = cli.Shell.AutosuggestModel
 		config.ShellAutosuggestTimeout = time.Duration(cli.Shell.AutosuggestTimeout) * time.Millisecond
 		config.ShellNewlineAutosuggestTimeout = time.Duration(cli.Shell.NewlineAutosuggestTimeout) * time.Millisecond
-		config.ShellColorDark = !cli.Shell.LightColor
+		config.ColorDark = !cli.LightColor
 		config.ShellMode = true
 		config.ShellLeavePromptAlone = cli.Shell.NoCommandPrompt
 		config.ShellMaxPromptTokens = cli.Shell.MaxPromptTokens

util/util.go

@@ -259,6 +259,7 @@ type StyleCodeblocksWriter struct {
 	terminalWidth int
 	normalColor   string
 	inlineColor   string
+	colorScheme   string
 	state         int
 	langSuffix    *bytes.Buffer
 	blockBuffer   *bytes.Buffer
@@ -270,17 +271,23 @@ func NewStyleCodeblocksWriter(
 	terminalWidth int,
 	normalColor string,
 	highlightColor string,
+	colorScheme string,
 ) *StyleCodeblocksWriter {
 	if terminalWidth == 0 {
 		panic("terminal width must be > 0")
 	}
 
+	if colorScheme == "" {
+		colorScheme = "monokai"
+	}
+
 	return &StyleCodeblocksWriter{
 		Writer:        writer,
 		state:         STATE_NEWLINE,
 		normalColor:   normalColor,
 		inlineColor:   highlightColor,
 		terminalWidth: terminalWidth,
+		colorScheme:   colorScheme,
 	}
 }
 
@@ -472,8 +479,13 @@ func lastLine(buff *bytes.Buffer, newlines int) []byte {
 func (this *StyleCodeblocksWriter) EndOfCodeLine(w io.Writer) error {
 	temp := new(bytes.Buffer)
 	blockBufferString := this.blockBuffer.String()
-	err := quick.Highlight(temp, blockBufferString,
-		this.langSuffix.String(), "terminal256", "monokai")
+
+	err := quick.Highlight(
+		temp,
+		blockBufferString,
+		this.langSuffix.String(),
+		"terminal256",
+		this.colorScheme)
 	if err != nil {
 		log.Printf("error highlighting code block: %s", err)
 	}

util/util_test.go

@@ -12,7 +12,7 @@ import (
 // Sanity test
 func TestStyleCodeblocksWriter(t *testing.T) {
 	buffer := new(bytes.Buffer)
-	writer := NewStyleCodeblocksWriter(buffer, 80, "", "")
+	writer := NewStyleCodeblocksWriter(buffer, 80, "", "", "")
 
 	writer.Write([]byte("Hello\n"))
 	writer.Write([]byte("```javascript\n"))
@@ -47,9 +47,13 @@ func TestStyleCodeblocksWriter(t *testing.T) {
 // 4. 1 backtick inlined, at start of line
 // 5. 3 backticks, indented
 
-func Test3BackticksNoLanguage(t *testing.T) {
+func getStyleCodeblocksWriter() (*bytes.Buffer, *StyleCodeblocksWriter) {
 	buffer := new(bytes.Buffer)
-	writer := NewStyleCodeblocksWriter(buffer, 80, "NORMAL", "HIGHLIGHT")
+	return buffer, NewStyleCodeblocksWriter(buffer, 80, "NORMAL", "HIGHLIGHT", "")
+}
+
+func Test3BackticksNoLanguage(t *testing.T) {
+	buffer, writer := getStyleCodeblocksWriter()
 
 	testStr := `Hello
 ` + "```" + `
@@ -67,8 +71,7 @@ Foo`
 }
 
 func Test3BackticksWithLanguage(t *testing.T) {
-	buffer := new(bytes.Buffer)
-	writer := NewStyleCodeblocksWriter(buffer, 80, "NORMAL", "HIGHLIGHT")
+	buffer, writer := getStyleCodeblocksWriter()
 
 	testStr := `Hello
 ` + "```javascript" + `
@@ -86,8 +89,7 @@ Foo`
 }
 
 func Test1BacktickInlined(t *testing.T) {
-	buffer := new(bytes.Buffer)
-	writer := NewStyleCodeblocksWriter(buffer, 80, "NORMAL", "HIGHLIGHT")
+	buffer, writer := getStyleCodeblocksWriter()
 
 	testStr := "Hello `console.log('Hi')` Foo"
 
@@ -100,8 +102,7 @@ func Test1BacktickInlined(t *testing.T) {
 }
 
 func Test1BacktickInlinedAtStartOfLine(t *testing.T) {
-	buffer := new(bytes.Buffer)
-	writer := NewStyleCodeblocksWriter(buffer, 80, "NORMAL", "HIGHLIGHT")
+	buffer, writer := getStyleCodeblocksWriter()
 
 	testStr := "`console.log('Hi')` Foo"
 
@@ -114,8 +115,7 @@ func Test1BacktickInlinedAtStartOfLine(t *testing.T) {
 }
 
 func Test3BackticksIndented(t *testing.T) {
-	buffer := new(bytes.Buffer)
-	writer := NewStyleCodeblocksWriter(buffer, 80, "NORMAL", "HIGHLIGHT")
+	buffer, writer := getStyleCodeblocksWriter()
 
 	testStr := `Hello