feat(agent): preserve fenced code highlighting

cmd/tui/chat/markdown.go

@@ -13,6 +13,17 @@ import (
 )
 
 func renderMarkdownForView(markdown string, width int) string {
+	return renderMarkdownForViewWithCodeCache(markdown, width, nil)
+}
+
+type markdownCodeBlockCacheKey struct {
+	language  string
+	source    string
+	formatter string
+	style     string
+}
+
+func renderMarkdownForViewWithCodeCache(markdown string, width int, codeCache *map[markdownCodeBlockCacheKey]string) string {
 	if width < 20 {
 		width = 20
 	}
@@ -21,21 +32,25 @@ func renderMarkdownForView(markdown string, width int) string {
 	var rendered []string
 	inCodeBlock := false
 	codeLanguage := ""
+	var codeLines []string
 	for i := 0; i < len(source); i++ {
 		line := strings.TrimRight(source[i], "\r")
 		trimmed := strings.TrimSpace(line)
 
 		if strings.HasPrefix(trimmed, "```") {
-			inCodeBlock = !inCodeBlock
 			if inCodeBlock {
-				codeLanguage = markdownCodeFenceLanguage(trimmed)
+				rendered = append(rendered, renderMarkdownCodeBlock(codeLines, codeLanguage, width, codeCache, true)...)
-			} else {
+				codeLines = nil
+				inCodeBlock = false
 				codeLanguage = ""
+			} else {
+				inCodeBlock = true
+				codeLanguage = markdownCodeFenceLanguage(trimmed)
 			}
 			continue
 		}
 		if inCodeBlock {
-			rendered = append(rendered, renderMarkdownCodeLine(line, codeLanguage, width)...)
+			codeLines = append(codeLines, line)
 			continue
 		}
 
@@ -56,6 +71,9 @@ func renderMarkdownForView(markdown string, width int) string {
 		}
 		rendered = append(rendered, wrapMarkdownInline(line, width)...)
 	}
+	if inCodeBlock {
+		rendered = append(rendered, renderMarkdownCodeBlock(codeLines, codeLanguage, width, codeCache, false)...)
+	}
 	return strings.Join(rendered, "\n")
 }
 
@@ -227,11 +245,12 @@ func renderMarkdownRunes(runes []markdownInlineRune) string {
 	return b.String()
 }
 
-func renderMarkdownCodeLine(line, language string, width int) []string {
+func renderMarkdownCodeBlock(source []string, language string, width int, codeCache *map[markdownCodeBlockCacheKey]string, complete bool) []string {
 	codeWidth := max(1, width-2)
-	lines := wrapChatText(line, codeWidth)
+	code := strings.Join(source, "\n")
-	if highlighted, ok := highlightMarkdownCodeLine(language, line); ok {
+	lines := wrapChatText(code, codeWidth)
-		lines = strings.Split(ansi.Hardwrap(highlighted, codeWidth, true), "\n")
+	if highlighted, ok := highlightMarkdownCodeBlock(language, code, codeCache, complete); ok {
+		lines = wrapHighlightedMarkdownCode(highlighted, codeWidth)
 	}
 	for i, wrapped := range lines {
 		lines[i] = "  " + chatCodeBlockStyle.Render(wrapped)
@@ -247,16 +266,81 @@ func markdownCodeFenceLanguage(fence string) string {
 	return strings.Fields(info)[0]
 }
 
-func highlightMarkdownCodeLine(language, line string) (string, bool) {
+func highlightMarkdownCodeBlock(language, source string, codeCache *map[markdownCodeBlockCacheKey]string, complete bool) (string, bool) {
-	if language == "" || line == "" || lipgloss.ColorProfile() == termenv.Ascii {
+	if language == "" || source == "" || lipgloss.ColorProfile() == termenv.Ascii {
-		return line, false
+		return source, false
+	}
+
+	key := markdownCodeBlockCacheKey{
+		language:  language,
+		source:    source,
+		formatter: markdownCodeFormatter(),
+		style:     markdownCodeStyle(),
+	}
+	if complete && codeCache != nil && *codeCache != nil {
+		if highlighted, ok := (*codeCache)[key]; ok {
+			return highlighted, true
+		}
 	}
 
 	var rendered strings.Builder
-	if err := quick.Highlight(&rendered, line, language, markdownCodeFormatter(), markdownCodeStyle()); err != nil {
+	if err := quick.Highlight(&rendered, source, language, key.formatter, key.style); err != nil {
-		return line, false
+		return source, false
 	}
-	return strings.TrimSuffix(rendered.String(), "\n"), true
+	highlighted := strings.TrimSuffix(rendered.String(), "\n")
+	if complete && codeCache != nil {
+		if *codeCache == nil {
+			*codeCache = make(map[markdownCodeBlockCacheKey]string)
+		}
+		(*codeCache)[key] = highlighted
+	}
+	return highlighted, true
+}
+
+func wrapHighlightedMarkdownCode(highlighted string, width int) []string {
+	lines := strings.Split(ansi.Hardwrap(highlighted, width, true), "\n")
+	activeStyle := ""
+	for i, line := range lines {
+		lines[i] = activeStyle + line
+		activeStyle = markdownCodeANSIStyle(activeStyle, line)
+		if activeStyle != "" && i < len(lines)-1 {
+			lines[i] += "\x1b[0m"
+		}
+	}
+	return lines
+}
+
+func markdownCodeANSIStyle(activeStyle, line string) string {
+	for {
+		start := strings.Index(line, "\x1b[")
+		if start < 0 {
+			return activeStyle
+		}
+		line = line[start+2:]
+		end := strings.IndexByte(line, 'm')
+		if end < 0 {
+			return activeStyle
+		}
+		sequence := line[:end]
+		if markdownCodeANSIReset(sequence) {
+			activeStyle = ""
+		} else {
+			activeStyle += "\x1b[" + sequence + "m"
+		}
+		line = line[end+1:]
+	}
+}
+
+func markdownCodeANSIReset(sequence string) bool {
+	if sequence == "" {
+		return true
+	}
+	for _, parameter := range strings.Split(sequence, ";") {
+		if parameter == "0" {
+			return true
+		}
+	}
+	return false
 }
 
 func markdownCodeFormatter() string {

cmd/tui/chat/render.go

@@ -37,6 +37,7 @@ type chatEntry struct {
 	version     int
 	renderKey   chatEntryRenderKey
 	renderLines []string
+	codeBlocks  map[markdownCodeBlockCacheKey]string
 }
 
 const (
@@ -202,7 +203,11 @@ func (m chatModel) renderEntryLinesCached(index int, entry chatEntry, body strin
 		}
 	}
 
-	lines := m.renderEntryLines(entry, body, width)
+	var codeCache *map[markdownCodeBlockCacheKey]string
+	if index >= 0 && index < len(m.entries) {
+		codeCache = &m.entries[index].codeBlocks
+	}
+	lines := m.renderEntryLines(entry, body, width, codeCache)
 	if index >= 0 && index < len(m.entries) {
 		m.entries[index].renderKey = key
 		m.entries[index].renderLines = slices.Clone(lines)
@@ -493,20 +498,20 @@ func (m chatModel) renderEntry(entry chatEntry) (string, string) {
 	}
 }
 
-func (m chatModel) renderEntryLines(entry chatEntry, body string, width int) []string {
+func (m chatModel) renderEntryLines(entry chatEntry, body string, width int, codeCache *map[markdownCodeBlockCacheKey]string) []string {
 	if width < 20 {
 		width = 20
 	}
 	switch entry.role {
 	case "assistant":
 		innerWidth := max(1, width-lipgloss.Width(chatMessageIndent))
-		lines := indentLines(splitRenderedBody(renderMarkdownForView(body, innerWidth)), chatMessageIndent)
+		lines := indentLines(splitRenderedBody(renderMarkdownForViewWithCodeCache(body, innerWidth, codeCache)), chatMessageIndent)
 		lines = append(lines, indentLines(renderMetricsLines(entry.metrics, innerWidth), chatMessageIndent)...)
 		return lines
 	case "thinking":
 		return renderThinkingLines(entry, width)
 	case "system", "slash":
-		return splitRenderedBody(renderMarkdownForView(body, width))
+		return splitRenderedBody(renderMarkdownForViewWithCodeCache(body, width, codeCache))
 	case "user":
 		return renderUserMessageLines(body, width)
 	case "compaction_summary":

cmd/tui/chat/render_test.go

@@ -2133,10 +2133,10 @@ func TestRenderMarkdownCodeFencesHighlightLanguageTaggedCode(t *testing.T) {
 			t.Fatalf("rendered code line width = %d, want <= 72: %q", got, stripANSI(line))
 		}
 	}
-	if got, ok := highlightMarkdownCodeLine("go", "package main"); !ok || !strings.Contains(got, "\x1b[") {
+	if got, ok := highlightMarkdownCodeBlock("go", "package main", nil, false); !ok || !strings.Contains(got, "\x1b[") {
 		t.Fatalf("language-tagged code should use Chroma highlighting: %q", got)
 	}
-	if got, ok := highlightMarkdownCodeLine("", "plain code stays plain"); ok || got != "plain code stays plain" {
+	if got, ok := highlightMarkdownCodeBlock("", "plain code stays plain", nil, false); ok || got != "plain code stays plain" {
 		t.Fatalf("untagged code = %q, want plain code", got)
 	}
 }
@@ -2146,7 +2146,88 @@ func TestRenderMarkdownCodeFencesRespectNoColorTerminals(t *testing.T) {
 	lipgloss.SetColorProfile(termenv.Ascii)
 	defer lipgloss.SetColorProfile(profile)
 
-	if got, ok := highlightMarkdownCodeLine("go", "package main"); ok || got != "package main" {
+	if got, ok := highlightMarkdownCodeBlock("go", "package main", nil, false); ok || got != "package main" {
 		t.Fatalf("no-color output = %q, highlighted = %t", got, ok)
 	}
 }
+
+func TestRenderMarkdownCodeFencesHighlightMultilineConstructs(t *testing.T) {
+	profile := lipgloss.ColorProfile()
+	lipgloss.SetColorProfile(termenv.TrueColor)
+	defer lipgloss.SetColorProfile(profile)
+
+	rendered := renderMarkdownForView(strings.Join([]string{
+		"```go",
+		"/*",
+		" * This is a multi-line comment.",
+		" */",
+		"```",
+	}, "\n"), 72)
+
+	for _, line := range strings.Split(rendered, "\n") {
+		if strings.Contains(stripANSI(line), "multi-line comment") && !strings.Contains(line, "\x1b[") {
+			t.Fatalf("multiline comment should remain highlighted: %q", line)
+		}
+	}
+}
+
+func TestWrapHighlightedMarkdownCodePreservesANSIStyle(t *testing.T) {
+	lines := wrapHighlightedMarkdownCode("\x1b[31mabcdef\x1b[0m", 3)
+	if got, want := len(lines), 2; got != want {
+		t.Fatalf("wrapped lines = %d, want %d: %#v", got, want, lines)
+	}
+	if got, want := stripANSI(lines[0]), "abc"; got != want {
+		t.Fatalf("first line = %q, want %q", got, want)
+	}
+	if got, want := stripANSI(lines[1]), "def"; got != want {
+		t.Fatalf("continuation = %q, want %q", got, want)
+	}
+	if !strings.HasSuffix(lines[0], "\x1b[0m") || !strings.HasPrefix(lines[1], "\x1b[31m") {
+		t.Fatalf("ANSI style should reset and resume at the wrap boundary: %#v", lines)
+	}
+}
+
+func TestHighlightMarkdownCodeBlockCachesCompletedBlocks(t *testing.T) {
+	profile := lipgloss.ColorProfile()
+	lipgloss.SetColorProfile(termenv.TrueColor)
+	defer lipgloss.SetColorProfile(profile)
+
+	cache := map[markdownCodeBlockCacheKey]string{}
+	if _, ok := highlightMarkdownCodeBlock("go", "package main", &cache, true); !ok || len(cache) != 1 {
+		t.Fatalf("completed code block should be highlighted and cached: %#v", cache)
+	}
+	for key := range cache {
+		cache[key] = "cached highlighting"
+	}
+	if got, ok := highlightMarkdownCodeBlock("go", "package main", &cache, true); !ok || got != "cached highlighting" {
+		t.Fatalf("completed code block should use cached highlighting: %q", got)
+	}
+}
+
+func TestChatCompletedCodeBlockSurvivesLaterStreamingRender(t *testing.T) {
+	profile := lipgloss.ColorProfile()
+	lipgloss.SetColorProfile(termenv.TrueColor)
+	defer lipgloss.SetColorProfile(profile)
+
+	m := chatModel{entries: []chatEntry{newChatEntry(chatEntry{
+		role: "assistant",
+		content: strings.Join([]string{
+			"```go",
+			"package main",
+			"```",
+		}, "\n"),
+	})}}
+	m.renderTranscript(80)
+	if got := len(m.entries[0].codeBlocks); got != 1 {
+		t.Fatalf("completed block cache entries = %d, want 1", got)
+	}
+	for key := range m.entries[0].codeBlocks {
+		m.entries[0].codeBlocks[key] = "\x1b[31mcached code\x1b[0m"
+	}
+
+	m.entries[0].content += "\ncontinued response"
+	m.markEntryDirty(0)
+	if got := stripANSI(m.renderTranscript(80)); !strings.Contains(got, "cached code") {
+		t.Fatalf("later streamed content should reuse completed block highlighting: %q", got)
+	}
+}