feat: add Set Prompt setting for prompt symbol

command.go

@@ -431,6 +431,8 @@ var Settings = map[string]CommandFunc{
 	"WaitPattern":   ExecuteSetWaitPattern,
 	"WaitTimeout":   ExecuteSetWaitTimeout,
 	"CursorBlink":   ExecuteSetCursorBlink,
+	"PromptColor":   ExecuteSetPromptColor,
+	"Prompt":        ExecuteSetPrompt,
 }
 
 // ExecuteSet applies the settings on the running vhs specified by the
@@ -700,6 +702,18 @@ func ExecuteSetCursorBlink(c parser.Command, v *VHS) error {
 	return nil
 }
 
+// ExecuteSetPromptColor sets the prompt color.
+func ExecuteSetPromptColor(c parser.Command, v *VHS) error {
+	v.Options.PromptColor = c.Args
+	return nil
+}
+
+// ExecuteSetPrompt sets the prompt symbol.
+func ExecuteSetPrompt(c parser.Command, v *VHS) error {
+	v.Options.Prompt = c.Args
+	return nil
+}
+
 // ExecuteScreenshot is a CommandFunc that indicates a new screenshot must be taken.
 func ExecuteScreenshot(c parser.Command, v *VHS) error {
 	v.ScreenshotNextFrame(c.Args)

evaluator.go

@@ -30,7 +30,7 @@ func Evaluate(ctx context.Context, tape string, out io.Writer, opts ...Evaluator
 
 	v := New()
 	for _, cmd := range cmds {
-		if cmd.Type == token.SET && cmd.Options == "Shell" || cmd.Type == token.ENV {
+		if cmd.Type == token.SET && (cmd.Options == "Shell" || cmd.Options == "PromptColor" || cmd.Options == "Prompt") || cmd.Type == token.ENV {
 			err := Execute(cmd, &v)
 			if err != nil {
 				return []error{err}
@@ -56,7 +56,7 @@ func Evaluate(ctx context.Context, tape string, out io.Writer, opts ...Evaluator
 	for i, cmd := range cmds {
 		if cmd.Type == token.SET || cmd.Type == token.OUTPUT || cmd.Type == token.REQUIRE {
 			_, _ = fmt.Fprintln(out, Highlight(cmd, false))
-			if cmd.Options != "Shell" {
+			if cmd.Options != "Shell" && cmd.Options != "PromptColor" && cmd.Options != "Prompt" {
 				err := Execute(cmd, &v)
 				if err != nil {
 					return []error{err}

shell.go

@@ -1,5 +1,10 @@
 package main
 
+import (
+	"fmt"
+	"strings"
+)
+
 // Supported shells of VH.
 const (
 	bash       = "bash"
@@ -13,64 +18,93 @@ const (
 	zsh        = "zsh"
 )
 
+// DefaultPromptColor is the default color for the shell prompt.
+const DefaultPromptColor = "#5B56E0"
+
+// DefaultPrompt is the default prompt symbol.
+const DefaultPrompt = ">"
+
 // Shell is a type that contains a prompt and the command to set up the shell.
 type Shell struct {
-	Command []string
-	Env     []string
+	Name string
 }
 
-// Shells contains a mapping from shell names to their Shell struct.
-var Shells = map[string]Shell{
-	bash: {
-		Env:     []string{"PS1=\\[\\e[38;2;90;86;224m\\]> \\[\\e[0m\\]", "BASH_SILENCE_DEPRECATION_WARNING=1"},
-		Command: []string{"bash", "--noprofile", "--norc", "--login", "+o", "history"},
-	},
-	zsh: {
-		Env:     []string{`PROMPT=%F{#5B56E0}> %F{reset_color}`},
-		Command: []string{"zsh", "--histnostore", "--no-rcs"},
-	},
-	fish: {
-		Command: []string{
+// ShellConfig returns the shell configuration with the given prompt and color.
+func ShellConfig(name, promptColor, prompt string) (env []string, command []string) {
+	// Parse hex color to RGB components
+	r, g, b := hexToRGB(promptColor)
+	hexNoHash := strings.TrimPrefix(promptColor, "#")
+
+	switch name {
+	case bash:
+		return []string{
+				fmt.Sprintf("PS1=\\[\\e[38;2;%d;%d;%dm\\]%s \\[\\e[0m\\]", r, g, b, prompt),
+				"BASH_SILENCE_DEPRECATION_WARNING=1",
+			},
+			[]string{"bash", "--noprofile", "--norc", "--login", "+o", "history"}
+	case zsh:
+		return []string{fmt.Sprintf(`PROMPT=%%F{#%s}%s %%F{reset_color}`, hexNoHash, prompt)},
+			[]string{"zsh", "--histnostore", "--no-rcs"}
+	case fish:
+		return nil, []string{
 			"fish",
 			"--login",
 			"--no-config",
 			"--private",
 			"-C", "function fish_greeting; end",
-			"-C", `function fish_prompt; set_color 5B56E0; echo -n "> "; set_color normal; end`,
-		},
-	},
-	powershell: {
-		Command: []string{
+			"-C", fmt.Sprintf(`function fish_prompt; set_color %s; echo -n "%s "; set_color normal; end`, hexNoHash, prompt),
+		}
+	case powershell:
+		return nil, []string{
 			"powershell",
 			"-NoLogo",
 			"-NoExit",
 			"-NoProfile",
 			"-Command",
-			`Set-PSReadLineOption -HistorySaveStyle SaveNothing; function prompt { Write-Host '>' -NoNewLine -ForegroundColor Blue; return ' ' }`,
-		},
-	},
-	pwsh: {
-		Command: []string{
+			fmt.Sprintf(`Set-PSReadLineOption -HistorySaveStyle SaveNothing; function prompt { Write-Host '%s' -NoNewLine -ForegroundColor ([System.Drawing.Color]::FromArgb(%d,%d,%d)); return ' ' }`, prompt, r, g, b),
+		}
+	case pwsh:
+		return nil, []string{
 			"pwsh",
 			"-Login",
 			"-NoLogo",
 			"-NoExit",
 			"-NoProfile",
 			"-Command",
-			`Set-PSReadLineOption -HistorySaveStyle SaveNothing; Function prompt { Write-Host -ForegroundColor Blue -NoNewLine '>'; return ' ' }`,
-		},
-	},
-	cmdexe: {
-		Command: []string{"cmd.exe", "/k", "prompt=^> "},
-	},
-	nushell: {
-		Command: []string{"nu", "--execute", "$env.PROMPT_COMMAND = {'\033[;38;2;91;86;224m>\033[m '}; $env.PROMPT_COMMAND_RIGHT = {''}"},
-	},
-	osh: {
-		Env:     []string{"PS1=\\[\\e[38;2;90;86;224m\\]> \\[\\e[0m\\]"},
-		Command: []string{"osh", "--norc"},
-	},
-	xonsh: {
-		Command: []string{"xonsh", "--no-rc", "-D", "PROMPT=\033[;38;2;91;86;224m>\033[m "},
-	},
+			fmt.Sprintf(`Set-PSReadLineOption -HistorySaveStyle SaveNothing; Function prompt { Write-Host -ForegroundColor ([System.Drawing.Color]::FromArgb(%d,%d,%d)) -NoNewLine '%s'; return ' ' }`, r, g, b, prompt),
+		}
+	case cmdexe:
+		return nil, []string{"cmd.exe", "/k", fmt.Sprintf("prompt=%s ", prompt)}
+	case nushell:
+		return nil, []string{"nu", "--execute", fmt.Sprintf("$env.PROMPT_COMMAND = {'\033[;38;2;%d;%d;%dm%s\033[m '}; $env.PROMPT_COMMAND_RIGHT = {''}", r, g, b, prompt)}
+	case osh:
+		return []string{fmt.Sprintf("PS1=\\[\\e[38;2;%d;%d;%dm\\]%s \\[\\e[0m\\]", r, g, b, prompt)},
+			[]string{"osh", "--norc"}
+	case xonsh:
+		return nil, []string{"xonsh", "--no-rc", "-D", fmt.Sprintf("PROMPT=\033[;38;2;%d;%d;%dm%s\033[m ", r, g, b, prompt)}
+	default:
+		return nil, nil
+	}
+}
+
+// hexToRGB converts a hex color string to RGB components.
+func hexToRGB(hex string) (r, g, b int) {
+	hex = strings.TrimPrefix(hex, "#")
+	if len(hex) == 6 {
+		fmt.Sscanf(hex, "%02x%02x%02x", &r, &g, &b)
+	}
+	return
+}
+
+// Shells contains a mapping from shell names to their Shell struct.
+var Shells = map[string]Shell{
+	bash:       {Name: bash},
+	zsh:        {Name: zsh},
+	fish:       {Name: fish},
+	powershell: {Name: powershell},
+	pwsh:       {Name: pwsh},
+	cmdexe:     {Name: cmdexe},
+	nushell:    {Name: nushell},
+	osh:        {Name: osh},
+	xonsh:      {Name: xonsh},
 }

shell_test.go

@@ -0,0 +1,104 @@
+package main
+
+import (
+	"strings"
+	"testing"
+)
+
+func TestHexToRGB(t *testing.T) {
+	tests := []struct {
+		hex     string
+		r, g, b int
+	}{
+		{"#5B56E0", 91, 86, 224},
+		{"5B56E0", 91, 86, 224},
+		{"#F6821F", 246, 130, 31},
+		{"#000000", 0, 0, 0},
+		{"#FFFFFF", 255, 255, 255},
+		{"", 0, 0, 0},
+		{"#FFF", 0, 0, 0}, // invalid length
+	}
+
+	for _, tc := range tests {
+		r, g, b := hexToRGB(tc.hex)
+		if r != tc.r || g != tc.g || b != tc.b {
+			t.Errorf("hexToRGB(%q) = (%d,%d,%d), want (%d,%d,%d)",
+				tc.hex, r, g, b, tc.r, tc.g, tc.b)
+		}
+	}
+}
+
+func TestShellConfigReturnsNonNil(t *testing.T) {
+	// Every known shell should return a non-nil command from ShellConfig.
+	shellNames := []string{
+		bash, zsh, fish, powershell, pwsh, cmdexe, nushell, osh, xonsh,
+	}
+
+	for _, name := range shellNames {
+		t.Run(name, func(t *testing.T) {
+			_, command := ShellConfig(name, DefaultPromptColor, DefaultPrompt)
+			if len(command) == 0 {
+				t.Errorf("ShellConfig(%q, %q) returned empty command", name, DefaultPromptColor)
+			}
+		})
+	}
+}
+
+func TestShellConfigCustomColor(t *testing.T) {
+	// Shells that use RGB values should embed the custom colour.
+	// bash uses PS1 with ANSI 38;2;R;G;B escape, so changing the colour
+	// should produce different RGB values in the env string.
+	env, _ := ShellConfig(bash, "#FF8000", DefaultPrompt)
+	if len(env) == 0 {
+		t.Fatal("expected env for bash")
+	}
+	// #FF8000 = 255,128,0
+	if !strings.Contains(env[0], "255;128;0") {
+		t.Errorf("bash PS1 does not contain expected RGB values for #FF8000: %s", env[0])
+	}
+
+	// zsh uses hex in the PROMPT string
+	env, _ = ShellConfig(zsh, "#FF8000", DefaultPrompt)
+	if len(env) == 0 {
+		t.Fatal("expected env for zsh")
+	}
+	if !strings.Contains(env[0], "FF8000") {
+		t.Errorf("zsh PROMPT does not contain hex colour FF8000: %s", env[0])
+	}
+}
+
+func TestShellConfigDefaultColor(t *testing.T) {
+	// With the default colour, bash should produce the original RGB values.
+	// #5B56E0 = 91,86,224
+	env, _ := ShellConfig(bash, DefaultPromptColor, DefaultPrompt)
+	if len(env) == 0 {
+		t.Fatal("expected env for bash")
+	}
+	if !strings.Contains(env[0], "91;86;224") {
+		t.Errorf("bash PS1 with default colour does not contain expected RGB: %s", env[0])
+	}
+}
+
+func TestShellConfigCustomPrompt(t *testing.T) {
+	// Every shell should embed the custom prompt symbol in its config.
+	shellNames := []string{
+		bash, zsh, fish, powershell, pwsh, cmdexe, nushell, osh, xonsh,
+	}
+
+	for _, name := range shellNames {
+		t.Run(name, func(t *testing.T) {
+			env, command := ShellConfig(name, DefaultPromptColor, "λ")
+			combined := strings.Join(env, " ") + " " + strings.Join(command, " ")
+			if !strings.Contains(combined, "λ") {
+				t.Errorf("ShellConfig(%q) with prompt λ does not contain the symbol: env=%v command=%v", name, env, command)
+			}
+		})
+	}
+}
+
+func TestShellConfigUnknownShell(t *testing.T) {
+	env, command := ShellConfig("unknown-shell", DefaultPromptColor, DefaultPrompt)
+	if env != nil || command != nil {
+		t.Errorf("expected nil for unknown shell, got env=%v command=%v", env, command)
+	}
+}

token/token.go

@@ -103,6 +103,8 @@ const (
 	WAIT_TIMEOUT    = "WAIT_TIMEOUT"    //nolint:revive
 	WAIT_PATTERN    = "WAIT_PATTERN"    //nolint:revive
 	CURSOR_BLINK    = "CURSOR_BLINK"    //nolint:revive
+	PROMPT_COLOR    = "PROMPT_COLOR"    //nolint:revive
+	PROMPT          = "PROMPT"         //nolint:revive
 )
 
 // Keywords maps keyword strings to tokens.
@@ -159,6 +161,8 @@ var Keywords = map[string]Type{
 	"Wait":          WAIT,
 	"Source":        SOURCE,
 	"CursorBlink":   CURSOR_BLINK,
+	"PromptColor":   PROMPT_COLOR,
+	"Prompt":        PROMPT,
 	"true":          BOOLEAN,
 	"false":         BOOLEAN,
 	"Screenshot":    SCREENSHOT,
@@ -173,7 +177,8 @@ func IsSetting(t Type) bool {
 	case SHELL, FONT_FAMILY, FONT_SIZE, LETTER_SPACING, LINE_HEIGHT,
 		FRAMERATE, TYPING_SPEED, THEME, PLAYBACK_SPEED, HEIGHT, WIDTH,
 		PADDING, LOOP_OFFSET, MARGIN_FILL, MARGIN, WINDOW_BAR,
-		WINDOW_BAR_SIZE, BORDER_RADIUS, CURSOR_BLINK, WAIT_TIMEOUT, WAIT_PATTERN:
+		WINDOW_BAR_SIZE, BORDER_RADIUS, CURSOR_BLINK, WAIT_TIMEOUT, WAIT_PATTERN,
+		PROMPT_COLOR, PROMPT:
 		return true
 	default:
 		return false

tty.go

@@ -24,7 +24,7 @@ func randomPort() int {
 }
 
 // buildTtyCmd builds the ttyd exec.Command on the given port.
-func buildTtyCmd(port int, shell Shell) *exec.Cmd {
+func buildTtyCmd(port int, shell Shell, promptColor, prompt string) *exec.Cmd {
 	args := []string{
 		fmt.Sprintf("--port=%d", port),
 		"--interface", "127.0.0.1",
@@ -36,11 +36,12 @@ func buildTtyCmd(port int, shell Shell) *exec.Cmd {
 		"--writable",
 	}
 
-	args = append(args, shell.Command...)
+	env, command := ShellConfig(shell.Name, promptColor, prompt)
+	args = append(args, command...)
 
 	cmd := exec.Command("ttyd", args...)
-	if shell.Env != nil {
-		cmd.Env = append(shell.Env, os.Environ()...)
+	if env != nil {
+		cmd.Env = append(env, os.Environ()...)
 	}
 	return cmd
 }

vhs.go

@@ -52,6 +52,8 @@ type Options struct {
 	CursorBlink   bool
 	Screenshot    ScreenshotOptions
 	Style         StyleOptions
+	PromptColor   string
+	Prompt        string
 }
 
 const (
@@ -107,6 +109,8 @@ func DefaultVHSOptions() Options {
 		Screenshot:    screenshot,
 		WaitTimeout:   defaultWaitTimeout,
 		WaitPattern:   defaultWaitPattern,
+		PromptColor:   DefaultPromptColor,
+		Prompt:        DefaultPrompt,
 	}
 }
 
@@ -131,7 +135,7 @@ func (vhs *VHS) Start() error {
 	}
 
 	port := randomPort()
-	vhs.tty = buildTtyCmd(port, vhs.Options.Shell)
+	vhs.tty = buildTtyCmd(port, vhs.Options.Shell, vhs.Options.PromptColor, vhs.Options.Prompt)
 	if err := vhs.tty.Start(); err != nil {
 		return fmt.Errorf("could not start tty: %w", err)
 	}