Vectrum
Sign In
288 lines · 9.2 KB · Go
package tools

import (
	"context"
	"fmt"
	"os"
	"strings"
)

// EditArgs is the edit_file input.
type EditArgs struct {
	Path       string `json:"path" jsonschema_description:"Path to the file, absolute or relative to the working root"`
	OldStr     string `json:"old_str" jsonschema_description:"Exact text to replace, including indentation. Must appear exactly once unless replace_all is set"`
	NewStr     string `json:"new_str" jsonschema_description:"Text to replace it with. Use an empty string to delete"`
	ReplaceAll bool   `json:"replace_all,omitempty" jsonschema_description:"Replace every occurrence instead of requiring a unique match"`
}

// NewEdit builds the edit_file tool. track must be the same tracker read_file
// writes to.
func NewEdit(root string, track FileTracker) Tool {
	return New("edit_file", strings.TrimSpace(`
Replace exact text in a file you have already read. old_str must match the file
byte for byte, including indentation, and must be unique — include surrounding
lines to disambiguate rather than hoping. Returns a unified diff of what
changed, which you should check against what you meant to do.`),
		func(_ context.Context, in EditArgs) (string, error) {
			path, err := resolve(root, in.Path)
			if err != nil {
				return "", err
			}
			if in.OldStr == "" {
				return "", fmt.Errorf("old_str is required; use write_file to create a new file")
			}
			if in.OldStr == in.NewStr {
				return "", fmt.Errorf("old_str and new_str are identical, so this edit would do nothing")
			}
			// Existence is checked before the read-before-edit rule, so a
			// wrong path reports the wrong path rather than "you have not read
			// this file", which would send the model looking for the wrong bug.
			raw, err := os.ReadFile(path)
			if err != nil {
				if os.IsNotExist(err) {
					return "", fmt.Errorf("no such file: %s", in.Path)
				}
				return "", err
			}
			if err := checkRead(track, path, in.Path); err != nil {
				return "", err
			}
			content, crlf := normalizeEndings(string(raw))
			oldText, newText := normalizeInput(in.OldStr), normalizeInput(in.NewStr)

			switch n := strings.Count(content, oldText); {
			case n == 0:
				return "", missReport(in.Path, content, oldText)
			case n > 1 && !in.ReplaceAll:
				return "", ambiguityReport(in.Path, content, oldText, n)
			}

			updated := content
			count := 1
			if in.ReplaceAll {
				count = strings.Count(content, oldText)
				updated = strings.ReplaceAll(content, oldText, newText)
			} else {
				updated = strings.Replace(content, oldText, newText, 1)
			}

			info, err := os.Stat(path)
			perm := os.FileMode(0o644)
			if err == nil {
				perm = info.Mode().Perm()
			}
			if err := os.WriteFile(path, []byte(restoreEndings(updated, crlf)), perm); err != nil {
				return "", err
			}
			// The file just changed, so the recorded read is stale by
			// definition. Re-point it at what is now on disk.
			if info, err := os.Stat(path); err == nil {
				track.MarkRead(path, info.ModTime())
			}

			return editSummary(in.Path, content, oldText, newText, count), nil
		})
}

// checkRead enforces the read-before-edit rule, and forces a re-read when the
// file has changed since. This single rule kills a large class of edits against
// remembered rather than observed content.
func checkRead(track FileTracker, path, display string) error {
	if track == nil {
		return nil
	}
	rec, ok := track.LastRead(path)
	if !ok {
		return fmt.Errorf("%s has not been read this session; read_file it first so you are editing what is actually there", display)
	}
	info, err := os.Stat(path)
	if err != nil {
		return nil // the read below will produce the real error
	}
	if info.ModTime().After(rec.Mod) {
		return fmt.Errorf("%s changed on disk since you read it; read_file it again before editing", display)
	}
	return nil
}

// missReport turns "not found" into something recoverable. A bare miss costs
// several turns of blind retrying; a miss carrying the actual text costs one.
func missReport(display, content, oldText string) error {
	oldLines := strings.Split(oldText, "\n")
	lines := strings.Split(content, "\n")

	// Whitespace is the usual culprit, so try again ignoring it.
	if start, end, ok := findNormalized(lines, oldLines); ok {
		var b strings.Builder
		fmt.Fprintf(&b, "old_str was not found in %s, but the same text is there with different whitespace at lines %d-%d. The file actually contains:\n\n",
			display, start+1, end)
		for i := start; i < end; i++ {
			fmt.Fprintf(&b, "%6d\t%s\n", i+1, capLine(lines[i]))
		}
		b.WriteString("\nCopy that text exactly, including its indentation.")
		return fmt.Errorf("%s", b.String())
	}

	// Otherwise show the neighbourhood of the closest line we can find.
	best := closestLine(lines, oldLines[0])
	if best < 0 {
		return fmt.Errorf("old_str was not found in %s, and nothing in the file resembles it. Read the file again — you may have the wrong path", display)
	}
	lo, hi := best-10, best+10
	if lo < 0 {
		lo = 0
	}
	if hi > len(lines) {
		hi = len(lines)
	}
	var b strings.Builder
	fmt.Fprintf(&b, "old_str was not found in %s. The closest thing in the file is around line %d:\n\n", display, best+1)
	for i := lo; i < hi; i++ {
		fmt.Fprintf(&b, "%6d\t%s\n", i+1, capLine(lines[i]))
	}
	b.WriteString("\nUse text copied from there, or read_file a wider range if the target is elsewhere.")
	return fmt.Errorf("%s", b.String())
}

// ambiguityReport refuses a non-unique match and says where every occurrence
// is, so the next attempt can be made unique. Silently taking the first match
// would corrupt the file in a way nobody notices for a week.
func ambiguityReport(display, content, oldText string, n int) error {
	var b strings.Builder
	fmt.Fprintf(&b, "old_str appears %d times in %s, at lines: %s.\n\n", n, display, joinInts(occurrenceLines(content, oldText)))
	b.WriteString("Add surrounding lines to old_str so it matches exactly one place, or set replace_all if every occurrence should change.")
	return fmt.Errorf("%s", b.String())
}

// editSummary renders the applied change as a unified diff.
func editSummary(display, content, oldText, newText string, count int) string {
	idx := strings.Index(content, oldText)
	startLine := strings.Count(content[:idx], "\n")

	lines := strings.Split(content, "\n")
	oldRegion := strings.Split(oldText, "\n")
	newRegion := strings.Split(newText, "\n")
	end := startLine + len(oldRegion)
	if end > len(lines) {
		end = len(lines)
	}
	before, after := contextAround(lines, startLine, end)

	out := unifiedDiff(display, before, oldRegion, newRegion, after, startLine+1)
	if count > 1 {
		return fmt.Sprintf("Replaced %d occurrences. First one:\n\n%s", count, out)
	}
	return out
}

// findNormalized looks for the target lines ignoring leading, trailing and
// repeated internal whitespace. Returns a half-open [start, end) line range.
func findNormalized(lines, target []string) (int, int, bool) {
	if len(target) == 0 || len(target) > len(lines) {
		return 0, 0, false
	}
	norm := make([]string, len(lines))
	for i, l := range lines {
		norm[i] = collapse(l)
	}
	want := make([]string, len(target))
	for i, l := range target {
		want[i] = collapse(l)
	}
	for i := 0; i+len(want) <= len(norm); i++ {
		match := true
		for j := range want {
			if norm[i+j] != want[j] {
				match = false
				break
			}
		}
		if match {
			return i, i + len(want), true
		}
	}
	return 0, 0, false
}

// closestLine scores every line against want and returns the best index, or -1
// when nothing is even loosely similar.
func closestLine(lines []string, want string) int {
	want = collapse(want)
	if want == "" {
		return -1
	}
	wantTokens := strings.Fields(want)
	best, bestScore := -1, 0.0
	for i, l := range lines {
		score := similarity(strings.Fields(collapse(l)), wantTokens)
		if score > bestScore {
			best, bestScore = i, score
		}
	}
	if bestScore < 0.34 {
		return -1
	}
	return best
}

// similarity is the fraction of want's tokens present in have.
func similarity(have, want []string) float64 {
	if len(want) == 0 {
		return 0
	}
	seen := make(map[string]int, len(have))
	for _, t := range have {
		seen[t]++
	}
	var hits int
	for _, t := range want {
		if seen[t] > 0 {
			seen[t]--
			hits++
		}
	}
	return float64(hits) / float64(len(want))
}

func collapse(s string) string { return strings.Join(strings.Fields(s), " ") }

// occurrenceLines returns the 1-indexed line of each occurrence of sub.
func occurrenceLines(content, sub string) []int {
	var out []int
	for off := 0; ; {
		i := strings.Index(content[off:], sub)
		if i < 0 {
			return out
		}
		out = append(out, strings.Count(content[:off+i], "\n")+1)
		off += i + len(sub)
	}
}

func joinInts(ns []int) string {
	parts := make([]string, len(ns))
	for i, n := range ns {
		parts[i] = fmt.Sprint(n)
	}
	return strings.Join(parts, ", ")
}

// normalizeEndings converts CRLF to LF and reports whether the file used CRLF,
// so the original ending can be restored on write.
func normalizeEndings(s string) (string, bool) {
	if strings.Contains(s, "\r\n") {
		return strings.ReplaceAll(s, "\r\n", "\n"), true
	}
	return s, false
}

func restoreEndings(s string, crlf bool) string {
	if !crlf {
		return s
	}
	return strings.ReplaceAll(s, "\n", "\r\n")
}

// normalizeInput strips CR from model-supplied text, which otherwise fails to
// match a file we have already normalized.
func normalizeInput(s string) string { return strings.ReplaceAll(s, "\r\n", "\n") }