fix(tui): escape ANSI escape sequences and control characters in tool

Ayman Bagabas created 4 months ago

call content

This would ensure that the content is displayed correctly in the
terminal, without any unintended formatting or control characters or
escape sequences. It will also style the escaped content for
display in the terminal, making it more readable.

Change summary

internal/tui/components/chat/messages/renderer.go | 41 +++++++++++++++++
1 file changed, 41 insertions(+)

Detailed changes

internal/tui/components/chat/messages/renderer.go

@@ -3,6 +3,7 @@ package messages
 import (
 	"encoding/json"
 	"fmt"
+	"strconv"
 	"strings"
 	"time"
 
@@ -656,6 +657,7 @@ func joinHeaderBody(header, body string) string {
 func renderPlainContent(v *toolCallCmp, content string) string {
 	t := styles.CurrentTheme()
 	content = strings.TrimSpace(content)
+	content = escapeContent(t, content)
 	lines := strings.Split(content, "\n")
 
 	width := v.textWidth() - 2 // -2 for left padding
@@ -694,6 +696,7 @@ func pad(v any, width int) string {
 
 func renderCodeContent(v *toolCallCmp, path, content string, offset int) string {
 	t := styles.CurrentTheme()
+	content = escapeContent(t, content)
 	truncated := truncateHeight(content, responseContextHeight)
 
 	highlighted, _ := highlight.SyntaxHighlight(truncated, path, t.BgBase)
@@ -766,3 +769,41 @@ func prettifyToolName(name string) string {
 		return name
 	}
 }
+
+// escapeContent escapes ANSI escape sequences and control characters in the
+// content and styles it for display in the terminal.
+func escapeContent(t *styles.Theme, content string) string {
+	lines := strings.Split(content, "\n")
+	for i, line := range lines {
+		lines[i] = escapeLine(t, line)
+	}
+
+	content = strings.Join(lines, "\n")
+	return content
+}
+
+// escapeLine escapes ANSI escape sequences and control characters and styles
+// them for display in the terminal.
+func escapeLine(t *styles.Theme, text string) string {
+	var (
+		sb    strings.Builder
+		state byte
+		seq   string
+		n     int
+		w     int
+	)
+	faint := ansi.NewStyle().Faint().ForegroundColor(t.FgMuted)
+	for len(text) > 0 {
+		seq, w, n, state = ansi.DecodeSequence(text, state, nil)
+		if w > 0 {
+			sb.WriteString(seq)
+		} else {
+			quote := strconv.Quote(seq)
+			quote = strings.TrimPrefix(quote, "\"")
+			quote = strings.TrimSuffix(quote, "\"")
+			sb.WriteString(faint.Styled(quote))
+		}
+		text = text[n:]
+	}
+	return sb.String()
+}