Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
178 changes: 167 additions & 11 deletions pkg/tui/components/tool/editfile/render.go
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down Expand Up @@ -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)

Expand All @@ -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")
}
Expand All @@ -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))
Expand All @@ -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:
Expand Down Expand Up @@ -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:]
}
Expand All @@ -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)
Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[HIGH] Tab expansion in prepareContent causes byte-offset mismatch with word-diff segments — wrong emphasis and potential panic on tab-indented lines

prepareContent() expands each \t to 4 spaces before passing content to syntaxHighlight. The chroma tokens produced therefore have byte offsets relative to the tab-expanded string. However, diffWords is called with the raw (non-tab-expanded) content:

// renderSplitSide — uses tab-expanded content for tokens:
content := prepareContent(line.Content)    // \t → "    " (4 bytes)
tokens := syntaxHighlight(content, filePath)

if len(wordSegs) > 0 {
    tokens = applyWordEmphasis(tokens, wordSegs)  // segments have unexpanded byte offsets!
}

// Word segs built with unexpanded content:
oldSegs, newSegs = diffWords(
    strings.TrimRight(pair.old.Content, "\r\n"),  // \t is still 1 byte here
    strings.TrimRight(pair.new.Content, "\r\n"),
)

applyWordEmphasis maps segment byte ranges onto chroma token byte ranges. For any line containing a \t before a changed word, the segment byte positions are shifted left by (tabWidth-1) * numTabsBefore relative to the chroma token byte positions. This causes:

  1. Wrong emphasis: Changed tokens are highlighted at the wrong positions (or not at all).
  2. Potential panic: The slice text[cursor-tokStart : end-tokStart] can produce a negative or out-of-bounds index if end (derived from the segment range) exceeds the actual chroma token's byte length.

Since Go source files — including the test fixtures in this PR — are almost always tab-indented, this affects every real-world usage of the word-level emphasis.

Fix: Call applyWordEmphasis using segments derived from prepareContent-processed input (i.e., pass prepareContent(...) to diffWords as well), or expand tabs in the content passed to diffWords before computing segments.

wrappedTokens := wrapTokens(tokens, width)

Expand All @@ -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)
}
Expand All @@ -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()
Expand Down
128 changes: 128 additions & 0 deletions pkg/tui/components/tool/editfile/render_test.go
Original file line number Diff line number Diff line change
@@ -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)
})
}
Loading
Loading