fix(ui): prevent AAAA probe bleed in terminals without Kitty graphics support (#1967)

M1xA and Ayman Bagabas created

* fix(ui): prevent AAAA probe bleed in terminals without Kitty graphics support

* refactor(ui): add an OS vendor type for Apple and use DRY for Kitty terminals

* refactor(ui): do not export private symbols, fix a LookupEnv for mocks

* refactor(ui): typo

* refactor(ui): remove dead code

* fix(ui): unify capability querying logic for terminal version and image capabilities

---------

Co-authored-by: Ayman Bagabas <ayman.bagabas@gmail.com>

Change summary

internal/cmd/root.go      |  17 +++-
internal/cmd/root_test.go | 160 +++++++++++++++++++++++++++++++++++++++++
internal/ui/model/ui.go   |  17 ++--
3 files changed, 181 insertions(+), 13 deletions(-)

Detailed changes

internal/cmd/root.go 🔗

@@ -33,6 +33,9 @@ import (
 	"github.com/spf13/cobra"
 )
 
+// kittyTerminals defines terminals supporting querying capabilities.
+var kittyTerminals = []string{"alacritty", "ghostty", "kitty", "rio", "wezterm"}
+
 func init() {
 	rootCmd.PersistentFlags().StringP("cwd", "c", "", "Current working directory")
 	rootCmd.PersistentFlags().StringP("data-dir", "D", "", "Custom crush data directory")
@@ -94,11 +97,11 @@ crush -y
 			slog.Info("New UI in control!")
 			com := common.DefaultCommon(app)
 			ui := ui.New(com)
-			ui.QueryVersion = shouldQueryTerminalVersion(env)
+			ui.QueryCapabilities = shouldQueryCapabilities(env)
 			model = ui
 		} else {
 			ui := tui.New(app)
-			ui.QueryVersion = shouldQueryTerminalVersion(env)
+			ui.QueryVersion = shouldQueryCapabilities(env)
 			model = ui
 		}
 		program := tea.NewProgram(
@@ -299,12 +302,16 @@ func createDotCrushDir(dir string) error {
 	return nil
 }
 
-func shouldQueryTerminalVersion(env uv.Environ) bool {
+func shouldQueryCapabilities(env uv.Environ) bool {
+	const osVendorTypeApple = "Apple"
 	termType := env.Getenv("TERM")
 	termProg, okTermProg := env.LookupEnv("TERM_PROGRAM")
 	_, okSSHTTY := env.LookupEnv("SSH_TTY")
+	if okTermProg && strings.Contains(termProg, osVendorTypeApple) {
+		return false
+	}
 	return (!okTermProg && !okSSHTTY) ||
-		(!strings.Contains(termProg, "Apple") && !okSSHTTY) ||
+		(!strings.Contains(termProg, osVendorTypeApple) && !okSSHTTY) ||
 		// Terminals that do support XTVERSION.
-		stringext.ContainsAny(termType, "alacritty", "ghostty", "kitty", "rio", "wezterm")
+		stringext.ContainsAny(termType, kittyTerminals...)
 }

internal/cmd/root_test.go 🔗

@@ -0,0 +1,160 @@
+package cmd
+
+import (
+	"strings"
+	"testing"
+
+	"github.com/charmbracelet/crush/internal/stringext"
+	uv "github.com/charmbracelet/ultraviolet"
+	"github.com/stretchr/testify/require"
+)
+
+type mockEnviron []string
+
+func (m mockEnviron) Getenv(key string) string {
+	v, _ := m.LookupEnv(key)
+	return v
+}
+
+func (m mockEnviron) LookupEnv(key string) (string, bool) {
+	for _, env := range m {
+		kv := strings.SplitN(env, "=", 2)
+		if len(kv) == 2 && kv[0] == key {
+			return kv[1], true
+		}
+	}
+	return "", false
+}
+
+func (m mockEnviron) ExpandEnv(s string) string {
+	return s // Not implemented for tests
+}
+
+func (m mockEnviron) Slice() []string {
+	return []string(m)
+}
+
+func TestShouldQueryImageCapabilities(t *testing.T) {
+	t.Parallel()
+
+	tests := []struct {
+		name string
+		env  mockEnviron
+		want bool
+	}{
+		{
+			name: "kitty terminal",
+			env:  mockEnviron{"TERM=xterm-kitty"},
+			want: true,
+		},
+		{
+			name: "wezterm terminal",
+			env:  mockEnviron{"TERM=xterm-256color"},
+			want: true,
+		},
+		{
+			name: "wezterm with WEZTERM env",
+			env:  mockEnviron{"TERM=xterm-256color", "WEZTERM_EXECUTABLE=/Applications/WezTerm.app/Contents/MacOS/wezterm-gui"},
+			want: true, // Not detected via TERM, only via stringext.ContainsAny which checks TERM
+		},
+		{
+			name: "Apple Terminal",
+			env:  mockEnviron{"TERM_PROGRAM=Apple_Terminal", "TERM=xterm-256color"},
+			want: false,
+		},
+		{
+			name: "alacritty",
+			env:  mockEnviron{"TERM=alacritty"},
+			want: true,
+		},
+		{
+			name: "ghostty",
+			env:  mockEnviron{"TERM=xterm-ghostty"},
+			want: true,
+		},
+		{
+			name: "rio",
+			env:  mockEnviron{"TERM=rio"},
+			want: true,
+		},
+		{
+			name: "wezterm (detected via TERM)",
+			env:  mockEnviron{"TERM=wezterm"},
+			want: true,
+		},
+		{
+			name: "SSH session",
+			env:  mockEnviron{"SSH_TTY=/dev/pts/0", "TERM=xterm-256color"},
+			want: false,
+		},
+		{
+			name: "generic terminal",
+			env:  mockEnviron{"TERM=xterm-256color"},
+			want: true,
+		},
+		{
+			name: "kitty over SSH",
+			env:  mockEnviron{"SSH_TTY=/dev/pts/0", "TERM=xterm-kitty"},
+			want: true,
+		},
+		{
+			name: "Apple Terminal with kitty TERM (should still be false due to TERM_PROGRAM)",
+			env:  mockEnviron{"TERM_PROGRAM=Apple_Terminal", "TERM=xterm-kitty"},
+			want: false,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			t.Parallel()
+			got := shouldQueryCapabilities(uv.Environ(tt.env))
+			require.Equal(t, tt.want, got, "shouldQueryImageCapabilities() = %v, want %v", got, tt.want)
+		})
+	}
+}
+
+// This is a helper to test the underlying logic of stringext.ContainsAny
+// which is used by shouldQueryImageCapabilities
+func TestStringextContainsAny(t *testing.T) {
+	t.Parallel()
+
+	tests := []struct {
+		name   string
+		s      string
+		substr []string
+		want   bool
+	}{
+		{
+			name:   "kitty in TERM",
+			s:      "xterm-kitty",
+			substr: kittyTerminals,
+			want:   true,
+		},
+		{
+			name:   "wezterm in TERM",
+			s:      "wezterm",
+			substr: kittyTerminals,
+			want:   true,
+		},
+		{
+			name:   "alacritty in TERM",
+			s:      "alacritty",
+			substr: kittyTerminals,
+			want:   true,
+		},
+		{
+			name:   "generic terminal not in list",
+			s:      "xterm-256color",
+			substr: kittyTerminals,
+			want:   false,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			t.Parallel()
+			got := stringext.ContainsAny(tt.s, tt.substr...)
+			require.Equal(t, tt.want, got)
+		})
+	}
+}

internal/ui/model/ui.go 🔗

@@ -145,9 +145,9 @@ type UI struct {
 	// terminal.
 	sendProgressBar bool
 
-	// QueryVersion instructs the TUI to query for the terminal version when it
+	// QueryCapabilities instructs the TUI to query for the terminal version when it
 	// starts.
-	QueryVersion bool
+	QueryCapabilities bool
 
 	// Editor components
 	textarea textarea.Model
@@ -300,7 +300,7 @@ func New(com *common.Common) *UI {
 // Init initializes the UI model.
 func (m *UI) Init() tea.Cmd {
 	var cmds []tea.Cmd
-	if m.QueryVersion {
+	if m.QueryCapabilities {
 		cmds = append(cmds, tea.RequestTerminalVersion)
 	}
 	// load the user commands async
@@ -351,11 +351,12 @@ func (m *UI) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
 			m.sendProgressBar = slices.Contains(msg, "WT_SESSION")
 		}
 		m.imgCaps.Env = uv.Environ(msg)
-		// XXX: Right now, we're using the same logic to determine image
-		// support. Terminals like Apple Terminal and possibly others might
-		// bleed characters when querying for Kitty graphics via APC escape
-		// sequences.
-		cmds = append(cmds, timage.RequestCapabilities(m.imgCaps.Env))
+		// Only query for image capabilities if the terminal is known to
+		// support Kitty graphics protocol. This prevents character bleeding
+		// on terminals that don't understand the APC escape sequences.
+		if m.QueryCapabilities {
+			cmds = append(cmds, timage.RequestCapabilities(m.imgCaps.Env))
+		}
 	case loadSessionMsg:
 		m.state = uiChat
 		if m.forceCompactMode {