Merge pull request #2900 from rumpl/improve-diff-view

feat(tui): word-level highlighting in edit_file diff view

Author: dgageot

pkg/tui/components/tool/editfile/render.go

@@ -71,6 +71,10 @@ func InvalidateCaches() {
 type chromaToken struct {
 	Text  string
 	Style lipgloss.Style
+	// Emphasized marks tokens whose underlying text was identified as part
+	// of a word-level diff. Renderers paint these with a stronger background
+	// so users can locate the precise edit within a long line.
+	Emphasized bool
 }
 
 type linePair struct {
@@ -310,18 +314,111 @@ func chromaToLipgloss(tokenType chroma.TokenType, style *chroma.Style) lipgloss.
 	return lipStyle
 }
 
+// segRange records the byte extent of a single word-diff segment within the
+// full line, plus whether that extent represents a change.
+type segRange struct {
+	start, end int
+	changed    bool
+}
+
+// applyWordEmphasis re-tags chroma tokens so that any portion of the token
+// text that falls inside a "changed" word-diff segment is split off into its
+// own emphasized token. Token boundaries from chroma and from the word-diff
+// rarely line up, so each chroma token is sliced against the segment cursor
+// to produce correctly emphasized sub-tokens.
+func applyWordEmphasis(tokens []chromaToken, segs []wordSegment) []chromaToken {
+	if len(segs) == 0 {
+		return tokens
+	}
+
+	ranges := make([]segRange, 0, len(segs))
+	pos := 0
+	for _, s := range segs {
+		ranges = append(ranges, segRange{start: pos, end: pos + len(s.Text), changed: s.Changed})
+		pos += len(s.Text)
+	}
+
+	out := make([]chromaToken, 0, len(tokens))
+	bytePos := 0
+	for _, tok := range tokens {
+		text := tok.Text
+		tokStart := bytePos
+		tokEnd := bytePos + len(text)
+
+		cursor := tokStart
+		for cursor < tokEnd {
+			r := findSegRange(ranges, cursor)
+			if r == nil {
+				out = append(out, chromaToken{
+					Text:  text[cursor-tokStart:],
+					Style: tok.Style,
+				})
+				break
+			}
+			end := min(tokEnd, r.end)
+			sub := text[cursor-tokStart : end-tokStart]
+			if sub != "" {
+				out = append(out, chromaToken{
+					Text:       sub,
+					Style:      tok.Style,
+					Emphasized: r.changed,
+				})
+			}
+			cursor = end
+		}
+		bytePos = tokEnd
+	}
+
+	return out
+}
+
+func findSegRange(ranges []segRange, pos int) *segRange {
+	for i := range ranges {
+		if pos >= ranges[i].start && pos < ranges[i].end {
+			return &ranges[i]
+		}
+	}
+	return nil
+}
+
+// emphasisStyleFor returns the per-side emphasis style (with a stronger
+// background tint) for the row's diff kind. Returns the unchanged style for
+// kinds that should never carry word-level emphasis.
+func emphasisStyleFor(kind udiff.OpKind) lipgloss.Style {
+	switch kind {
+	case udiff.Delete:
+		return styles.DiffRemoveEmphStyle
+	case udiff.Insert:
+		return styles.DiffAddEmphStyle
+	default:
+		return styles.DiffUnchangedStyle
+	}
+}
+
 func renderDiffWithSyntaxHighlight(diff []*udiff.Hunk, filePath string, width int) string {
 	var output strings.Builder
 	contentWidth := width - lineNumWidth
 
 	for _, hunk := range diff {
+		// Build word-diff lookups for paired delete/insert lines so we can
+		// emphasize the precise tokens that changed.
+		wordDiffs := buildLineWordDiffs(hunk.Lines)
+
 		oldLineNum := hunk.FromLine
 		newLineNum := hunk.ToLine
 
-		for _, line := range hunk.Lines {
+		for li, line := range hunk.Lines {
 			lineNum := getDisplayLineNumber(&line, &oldLineNum, &newLineNum)
 			content := prepareContent(line.Content)
 			tokens := syntaxHighlight(content, filePath)
+			if wd, ok := wordDiffs[li]; ok {
+				switch line.Kind {
+				case udiff.Delete:
+					tokens = applyWordEmphasis(tokens, wd.old)
+				case udiff.Insert:
+					tokens = applyWordEmphasis(tokens, wd.new)
+				}
+			}
 			lineStyle := getLineStyle(line.Kind)
 			wrappedTokens := wrapTokens(tokens, contentWidth)
 
@@ -334,7 +431,7 @@ func renderDiffWithSyntaxHighlight(diff []*udiff.Hunk, filePath string, width in
 					// Use continuation indicator for wrapped lines
 					lineNumStr = styles.LineNumberStyle.Render("   → ")
 				}
-				rendered := renderTokensWithStyle(tokenLine, lineStyle)
+				rendered := renderTokensWithStyle(tokenLine, lineStyle, line.Kind)
 				padded := padToWidth(rendered, contentWidth, lineStyle)
 				output.WriteString(lineNumStr + padded + "\n")
 			}
@@ -358,8 +455,22 @@ func renderSplitDiffWithSyntaxHighlight(diff []*udiff.Hunk, filePath string, wid
 
 	for _, hunk := range diff {
 		for _, pair := range pairDiffLines(hunk.Lines, hunk.FromLine, hunk.ToLine) {
-			leftLines := renderSplitSide(pair.old, pair.oldLineNum, filePath, contentWidth)
-			rightLines := renderSplitSide(pair.new, pair.newLineNum, filePath, contentWidth)
+			// Word-diff is only meaningful when both halves are present
+			// and represent a delete/insert pair. Inputs go through
+			// prepareContent so the segment byte offsets line up with the
+			// chroma tokens produced for rendering (which also receive
+			// tab-expanded content).
+			var oldSegs, newSegs []wordSegment
+			if pair.old != nil && pair.new != nil &&
+				pair.old.Kind == udiff.Delete && pair.new.Kind == udiff.Insert {
+				oldSegs, newSegs = diffWords(
+					prepareContent(pair.old.Content),
+					prepareContent(pair.new.Content),
+				)
+			}
+
+			leftLines := renderSplitSide(pair.old, pair.oldLineNum, filePath, contentWidth, oldSegs)
+			rightLines := renderSplitSide(pair.new, pair.newLineNum, filePath, contentWidth, newSegs)
 
 			// Ensure both sides have the same number of lines for alignment
 			maxLines := max(len(rightLines), len(leftLines))
@@ -383,6 +494,32 @@ func renderSplitDiffWithSyntaxHighlight(diff []*udiff.Hunk, filePath string, wid
 	return strings.TrimSuffix(output.String(), "\n")
 }
 
+// lineWordDiff holds the per-side segment arrays computed for one
+// delete/insert pair. It is keyed by the *delete* line index — the matching
+// insert sits directly after it.
+type lineWordDiff struct {
+	old []wordSegment
+	new []wordSegment
+}
+
+func buildLineWordDiffs(lines []udiff.Line) map[int]lineWordDiff {
+	out := map[int]lineWordDiff{}
+	for i := range len(lines) - 1 {
+		if lines[i].Kind != udiff.Delete || lines[i+1].Kind != udiff.Insert {
+			continue
+		}
+		// Use prepareContent so segment offsets align with the chroma
+		// tokens (which are produced from the same tab-expanded text).
+		oldText := prepareContent(lines[i].Content)
+		newText := prepareContent(lines[i+1].Content)
+		oldSegs, newSegs := diffWords(oldText, newText)
+		wd := lineWordDiff{old: oldSegs, new: newSegs}
+		out[i] = wd
+		out[i+1] = wd
+	}
+	return out
+}
+
 func getDisplayLineNumber(line *udiff.Line, oldLineNum, newLineNum *int) int {
 	switch line.Kind {
 	case udiff.Delete:
@@ -456,7 +593,11 @@ func wrapTokens(tokens []chromaToken, maxWidth int) [][]chromaToken {
 				fitWidth = runewidth.RuneWidth(r)
 			}
 
-			currentLine = append(currentLine, chromaToken{Text: text[:fitLen], Style: token.Style})
+			currentLine = append(currentLine, chromaToken{
+				Text:       text[:fitLen],
+				Style:      token.Style,
+				Emphasized: token.Emphasized,
+			})
 			currentWidth += fitWidth
 			text = text[fitLen:]
 		}
@@ -473,14 +614,20 @@ func wrapTokens(tokens []chromaToken, maxWidth int) [][]chromaToken {
 	return lines
 }
 
-// renderSplitSide renders a split side with text wrapping support
-func renderSplitSide(line *udiff.Line, lineNum int, filePath string, width int) []string {
+// renderSplitSide renders a split side with text wrapping support.
+// When wordSegs is non-nil and the line is part of a delete/insert pair, the
+// chroma tokens are re-tagged so the changed substrings render with a
+// stronger background tint.
+func renderSplitSide(line *udiff.Line, lineNum int, filePath string, width int, wordSegs []wordSegment) []string {
 	if line == nil {
 		return []string{renderEmptySplitSide(width)}
 	}
 
 	content := prepareContent(line.Content)
 	tokens := syntaxHighlight(content, filePath)
+	if len(wordSegs) > 0 {
+		tokens = applyWordEmphasis(tokens, wordSegs)
+	}
 	lineStyle := getLineStyle(line.Kind)
 	wrappedTokens := wrapTokens(tokens, width)
 
@@ -494,7 +641,7 @@ func renderSplitSide(line *udiff.Line, lineNum int, filePath string, width int)
 			// Use continuation indicator for wrapped lines
 			lineNumStr = "   → "
 		}
-		rendered := renderTokensWithStyle(tokenLine, lineStyle)
+		rendered := renderTokensWithStyle(tokenLine, lineStyle, line.Kind)
 		padded := padToWidth(rendered, width, lineStyle)
 		result = append(result, styles.LineNumberStyle.Render(lineNumStr)+padded)
 	}
@@ -509,12 +656,21 @@ func renderEmptySplitSide(width int) string {
 	return styles.LineNumberStyle.Render(lineNumStr) + emptySpace
 }
 
-func renderTokensWithStyle(tokens []chromaToken, lineStyle lipgloss.Style) string {
+func renderTokensWithStyle(tokens []chromaToken, lineStyle lipgloss.Style, kind udiff.OpKind) string {
 	var output strings.Builder
 
+	emph := emphasisStyleFor(kind)
 	for _, token := range tokens {
-		styledToken := token.Style.Background(lineStyle.GetBackground())
-		output.WriteString(styledToken.Render(token.Text))
+		if token.Emphasized && (kind == udiff.Delete || kind == udiff.Insert) {
+			// Keep the chroma foreground so syntax colors carry through into
+			// the emphasized block, but override the background with the
+			// stronger emphasis tint and add bold for extra weight.
+			style := token.Style.Background(emph.GetBackground()).Bold(true)
+			output.WriteString(style.Render(token.Text))
+			continue
+		}
+		style := token.Style.Background(lineStyle.GetBackground())
+		output.WriteString(style.Render(token.Text))
 	}
 
 	return output.String()

pkg/tui/components/tool/editfile/render_test.go

@@ -0,0 +1,128 @@
+package editfile
+
+import (
+	"encoding/json"
+	"os"
+	"path/filepath"
+	"strings"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	"github.com/docker/docker-agent/pkg/tools"
+	"github.com/docker/docker-agent/pkg/tui/types"
+)
+
+// TestRenderEditFile_EndToEnd writes a temporary source file, builds an
+// edit_file tool call against it, and renders both unified and split views.
+// The test focuses on structural elements rather than exact escape sequences,
+// which depend on the active theme.
+func TestRenderEditFile_EndToEnd(t *testing.T) {
+	dir := t.TempDir()
+	path := filepath.Join(dir, "main.go")
+
+	updated := `package main
+
+import "fmt"
+
+func main() {
+	x := 10
+	y := 20
+	fmt.Println(x + y)
+}
+`
+	// Simulate the post-execution state: file already contains the new content.
+	require.NoError(t, os.WriteFile(path, []byte(updated), 0o644))
+
+	args := map[string]any{
+		"path": path,
+		"edits": []map[string]string{
+			{
+				"oldText": "\tx := 1\n\ty := 2",
+				"newText": "\tx := 10\n\ty := 20",
+			},
+		},
+	}
+	encoded, err := json.Marshal(args)
+	require.NoError(t, err)
+
+	toolCall := tools.ToolCall{
+		ID: "test-render-1",
+		Function: tools.FunctionCall{
+			Name:      "edit_file",
+			Arguments: string(encoded),
+		},
+	}
+
+	// Reset cache so the test is hermetic across runs.
+	InvalidateCaches()
+	t.Cleanup(InvalidateCaches)
+
+	unified := renderEditFile(toolCall, 120, false, types.ToolStatusCompleted)
+	split := renderEditFile(toolCall, 120, true, types.ToolStatusCompleted)
+
+	for _, out := range []string{unified, split} {
+		assert.NotEmpty(t, out)
+		// Source content should appear in the diff regardless of theme escapes.
+		assert.True(t, strings.Contains(out, "10") || strings.Contains(out, "20"))
+	}
+
+	added, removed := countDiffLines(toolCall, types.ToolStatusCompleted)
+	assert.Equal(t, 2, added)
+	assert.Equal(t, 2, removed)
+}
+
+func TestRenderEditFile_TabIndentedLineDoesNotPanic(t *testing.T) {
+	// Regression: tab-indented modified lines used to feed raw (1-byte-tab)
+	// text into diffWords while chroma tokens were built from the
+	// tab-expanded variant, producing out-of-bounds slice indices in
+	// applyWordEmphasis. The fix routes both through prepareContent.
+	dir := t.TempDir()
+	path := filepath.Join(dir, "main.go")
+
+	updated := "package main\n\nfunc main() {\n\tx := 10\n}\n"
+	require.NoError(t, os.WriteFile(path, []byte(updated), 0o644))
+
+	args := map[string]any{
+		"path": path,
+		"edits": []map[string]string{
+			{"oldText": "\tx := 1", "newText": "\tx := 10"},
+		},
+	}
+	encoded, _ := json.Marshal(args)
+	toolCall := tools.ToolCall{
+		ID:       "test-tab-1",
+		Function: tools.FunctionCall{Name: "edit_file", Arguments: string(encoded)},
+	}
+
+	InvalidateCaches()
+	t.Cleanup(InvalidateCaches)
+
+	assert.NotPanics(t, func() {
+		_ = renderEditFile(toolCall, 120, false, types.ToolStatusCompleted)
+		_ = renderEditFile(toolCall, 120, true, types.ToolStatusCompleted)
+	})
+}
+
+func TestRenderEditFile_MissingFileReturnsEmptyDiff(t *testing.T) {
+	args := map[string]any{
+		"path": "/nonexistent/path/that/does/not/exist.go",
+		"edits": []map[string]string{
+			{"oldText": "a", "newText": "b"},
+		},
+	}
+	encoded, _ := json.Marshal(args)
+	toolCall := tools.ToolCall{
+		ID:       "test-missing-1",
+		Function: tools.FunctionCall{Name: "edit_file", Arguments: string(encoded)},
+	}
+
+	InvalidateCaches()
+	t.Cleanup(InvalidateCaches)
+
+	// Should not panic on a missing source file.
+	assert.NotPanics(t, func() {
+		_ = renderEditFile(toolCall, 100, false, types.ToolStatusCompleted)
+	})
+}