diff --git a/pkg/tools/shared/diff_result.go b/pkg/tools/shared/diff_result.go index 99f7e9793..0f6166219 100644 --- a/pkg/tools/shared/diff_result.go +++ b/pkg/tools/shared/diff_result.go @@ -1,6 +1,7 @@ package toolshared import ( + "bytes" "fmt" "path/filepath" "strings" @@ -8,11 +9,15 @@ import ( "github.com/pmezard/go-difflib/difflib" ) -const noContentChangeDiffMessage = "(no content change)" +const ( + noContentChangeDiffMessage = "(no content change)" + noNewlineAtEOFMarker = `\ No newline at end of file` +) // DiffResult creates a user-visible tool result containing a unified diff for // a successful file edit. The diff is included for both the LLM and the user so -// the follow-up assistant response can reason about the exact change set. +// the follow-up assistant response can reason about the resulting change set, +// including EOF newline transitions. func DiffResult(path string, before, after []byte) *ToolResult { diff, err := buildUnifiedDiff(path, before, after) if err != nil { @@ -25,8 +30,8 @@ func DiffResult(path string, before, after []byte) *ToolResult { func buildUnifiedDiff(path string, before, after []byte) (string, error) { diff, err := difflib.GetUnifiedDiffString(difflib.UnifiedDiff{ - A: difflib.SplitLines(string(before)), - B: difflib.SplitLines(string(after)), + A: splitDiffLinesPreservingEOF(before), + B: splitDiffLinesPreservingEOF(after), FromFile: "a/" + diffDisplayPath(path), ToFile: "b/" + diffDisplayPath(path), Context: 3, @@ -43,6 +48,36 @@ func buildUnifiedDiff(path string, before, after []byte) (string, error) { return diff, nil } +func splitDiffLinesPreservingEOF(content []byte) []string { + if len(content) == 0 { + return nil + } + + lines := make([]string, 0, bytes.Count(content, []byte{'\n'})+1) + lineStart := 0 + for i, b := range content { + if b != '\n' { + continue + } + lines = append(lines, string(content[lineStart:i+1])) + lineStart = i + 1 + } + if lineStart < len(content) { + lines = append(lines, string(content[lineStart:])) + } + + if lacksTrailingNewline(content) { + lines[len(lines)-1] += "\n" + lines = append(lines, noNewlineAtEOFMarker+"\n") + } + + return lines +} + +func lacksTrailingNewline(content []byte) bool { + return len(content) > 0 && !bytes.HasSuffix(content, []byte("\n")) +} + func diffDisplayPath(path string) string { displayPath := strings.TrimLeft(filepath.ToSlash(path), "/") if displayPath == "" { diff --git a/pkg/tools/shared/diff_result_test.go b/pkg/tools/shared/diff_result_test.go index 7c06d205a..848bde221 100644 --- a/pkg/tools/shared/diff_result_test.go +++ b/pkg/tools/shared/diff_result_test.go @@ -26,7 +26,7 @@ func TestDiffResult_UserVisibleUnifiedDiff(t *testing.T) { "```diff", "--- a/tmp/example.txt", "+++ b/tmp/example.txt", - "@@ -1,4 +1,4 @@", + "@@ -1,3 +1,3 @@", " alpha", "-beta", "+beta 2", @@ -48,6 +48,42 @@ func TestBuildUnifiedDiff_NoContentChange(t *testing.T) { } } +func TestBuildUnifiedDiff_PreservesTrailingNewlineRemoval(t *testing.T) { + diff, err := buildUnifiedDiff("test.txt", []byte("same\n"), []byte("same")) + if err != nil { + t.Fatalf("buildUnifiedDiff() error = %v", err) + } + + for _, want := range []string{ + "--- a/test.txt", + "+++ b/test.txt", + " same", + "+" + noNewlineAtEOFMarker, + } { + if !strings.Contains(diff, want) { + t.Fatalf("buildUnifiedDiff() missing %q:\n%s", want, diff) + } + } +} + +func TestBuildUnifiedDiff_PreservesTrailingNewlineAddition(t *testing.T) { + diff, err := buildUnifiedDiff("test.txt", []byte("same"), []byte("same\n")) + if err != nil { + t.Fatalf("buildUnifiedDiff() error = %v", err) + } + + for _, want := range []string{ + "--- a/test.txt", + "+++ b/test.txt", + " same", + "-" + noNewlineAtEOFMarker, + } { + if !strings.Contains(diff, want) { + t.Fatalf("buildUnifiedDiff() missing %q:\n%s", want, diff) + } + } +} + func TestBuildUnifiedDiff_UsesNormalizedDisplayPaths(t *testing.T) { diff, err := buildUnifiedDiff("/tmp/nested/example.txt", []byte("before\n"), []byte("after\n")) if err != nil {