claudetool/patch.go - sketch - Gitiles

 package claudetool

 import (
 	"bytes"
 	"context"
 	"encoding/json"
 	"errors"
 	"fmt"
 	"go/parser"
 	"go/token"
 	"log/slog"
 	"os"
 	"path/filepath"
 	"strings"

 	"sketch.dev/ant"
 	"sketch.dev/claudetool/editbuf"
 	"sketch.dev/claudetool/patchkit"
 )

 // Patch is a tool for precise text modifications in files.
 var Patch = &ant.Tool{
 	Name:        PatchName,
 	Description: strings.TrimSpace(PatchDescription),
 	InputSchema: ant.MustSchema(PatchInputSchema),
 	Run:         PatchRun,
 }

 const (
 	PatchName        = "patch"
 	PatchDescription = `
 File modification tool for precise text edits.

 Operations:
 - replace: Substitute text with new content
 - append_eof: Append new text at the end of the file
 - prepend_bof: Insert new text at the beginning of the file
 - overwrite: Replace the entire file with new content (automatically creates the file)

 Usage notes:
 - All inputs are interpreted literally (no automatic newline or whitespace handling)
 - For replace operations, oldText must appear EXACTLY ONCE in the file
 `

 	// If you modify this, update the termui template for prettier rendering.
 	PatchInputSchema = `
 {
   "type": "object",
   "required": ["path", "patches"],
   "properties": {
     "path": {
       "type": "string",
       "description": "Absolute path to the file to patch"
     },
     "patches": {
       "type": "array",
       "description": "List of patch requests to apply",
       "items": {
         "type": "object",
         "required": ["operation", "newText"],
         "properties": {
           "operation": {
             "type": "string",
             "enum": ["replace", "append_eof", "prepend_bof", "overwrite"],
             "description": "Type of operation to perform"
           },
           "oldText": {
             "type": "string",
             "description": "Text to locate for the operation (must be unique in file, required for replace)"
           },
           "newText": {
             "type": "string",
             "description": "The new text to use (empty for deletions)"
           }
         }
       }
     }
   }
 }
 `
 )

 // TODO: maybe rename PatchRequest to PatchOperation or PatchSpec or PatchPart or just Patch?

 type patchInput struct {
 	Path    string         `json:"path"`
 	Patches []patchRequest `json:"patches"`
 }

 type patchRequest struct {
 	Operation string `json:"operation"`
 	OldText   string `json:"oldText,omitempty"`
 	NewText   string `json:"newText,omitempty"`
 }

 // PatchRun is the entry point for the user_patch tool.
 func PatchRun(ctx context.Context, m json.RawMessage) (string, error) {
 	var input patchInput
 	if err := json.Unmarshal(m, &input); err != nil {
 		return "", fmt.Errorf("failed to unmarshal user_patch input: %w", err)
 	}

 	// Validate the input
 	if !filepath.IsAbs(input.Path) {
 		return "", fmt.Errorf("path %q is not absolute", input.Path)
 	}
 	if len(input.Patches) == 0 {
 		return "", fmt.Errorf("no patches provided")
 	}
 	// TODO: check whether the file is autogenerated, and if so, require a "force" flag to modify it.

 	orig, err := os.ReadFile(input.Path)
 	// If the file doesn't exist, we can still apply patches
 	// that don't require finding existing text.
 	switch {
 	case errors.Is(err, os.ErrNotExist):
 		for _, patch := range input.Patches {
 			switch patch.Operation {
 			case "prepend_bof", "append_eof", "overwrite":
 			default:
 				return "", fmt.Errorf("file %q does not exist", input.Path)
 			}
 		}
 	case err != nil:
 		return "", fmt.Errorf("failed to read file %q: %w", input.Path, err)
 	}

 	likelyGoFile := strings.HasSuffix(input.Path, ".go")

 	autogenerated := likelyGoFile && isAutogeneratedGoFile(orig)
 	parsed := likelyGoFile && parseGo(orig) != nil

 	origStr := string(orig)
 	// Process the patches "simultaneously", minimizing them along the way.
 	// Claude generates patches that interact with each other.
 	buf := editbuf.NewBuffer(orig)

 	// TODO: is it better to apply the patches that apply cleanly and report on the failures?
 	// or instead have it be all-or-nothing?
 	// For now, it is all-or-nothing.
 	// TODO: when the model gets into a "cannot apply patch" cycle of doom, how do we get it unstuck?
 	// Also: how do we detect that it's in a cycle?
 	var patchErr error
 	for i, patch := range input.Patches {
 		switch patch.Operation {
 		case "prepend_bof":
 			buf.Insert(0, patch.NewText)
 		case "append_eof":
 			buf.Insert(len(orig), patch.NewText)
 		case "overwrite":
 			buf.Replace(0, len(orig), patch.NewText)
 		case "replace":
 			if patch.OldText == "" {
 				return "", fmt.Errorf("patch %d: oldText cannot be empty for %s operation", i, patch.Operation)
 			}

 			// Attempt to apply the patch.
 			spec, count := patchkit.Unique(origStr, patch.OldText, patch.NewText)
 			switch count {
 			case 0:
 				// no matches, maybe recoverable, continued below
 			case 1:
 				// exact match, apply
 				slog.DebugContext(ctx, "patch_applied", "method", "unique")
 				spec.ApplyToEditBuf(buf)
 				continue
 			case 2:
 				// multiple matches
 				patchErr = errors.Join(patchErr, fmt.Errorf("old text not unique:\n%s", patch.OldText))
 			default:
 				// TODO: return an error instead of using agentPatch
 				slog.ErrorContext(ctx, "unique returned unexpected count", "count", count)
 				patchErr = errors.Join(patchErr, fmt.Errorf("internal error"))
 				continue
 			}

 			// The following recovery mechanisms are heuristic.
 			// They aren't perfect, but they appear safe,
 			// and the cases they cover appear with some regularity.

 			// Try adjusting the whitespace prefix.
 			spec, ok := patchkit.UniqueDedent(origStr, patch.OldText, patch.NewText)
 			if ok {
 				slog.DebugContext(ctx, "patch_applied", "method", "unique_dedent")
 				spec.ApplyToEditBuf(buf)
 				continue
 			}

 			// Try ignoring leading/trailing whitespace in a semantically safe way.
 			spec, ok = patchkit.UniqueInValidGo(origStr, patch.OldText, patch.NewText)
 			if ok {
 				slog.DebugContext(ctx, "patch_applied", "method", "unique_in_valid_go")
 				spec.ApplyToEditBuf(buf)
 				continue
 			}

 			// Try ignoring semantically insignificant whitespace.
 			spec, ok = patchkit.UniqueGoTokens(origStr, patch.OldText, patch.NewText)
 			if ok {
 				slog.DebugContext(ctx, "patch_applied", "method", "unique_go_tokens")
 				spec.ApplyToEditBuf(buf)
 				continue
 			}

 			// Try trimming the first line of the patch, if we can do so safely.
 			spec, ok = patchkit.UniqueTrim(origStr, patch.OldText, patch.NewText)
 			if ok {
 				slog.DebugContext(ctx, "patch_applied", "method", "unique_trim")
 				spec.ApplyToEditBuf(buf)
 				continue
 			}

 			// No dice.
 			patchErr = errors.Join(patchErr, fmt.Errorf("old text not found:\n%s", patch.OldText))
 			continue
 		default:
 			return "", fmt.Errorf("unrecognized operation %q", patch.Operation)
 		}
 	}

 	if patchErr != nil {
 		sendTelemetry(ctx, "patch_error", map[string]any{
 			"orig":    origStr,
 			"patches": input.Patches,
 			"errors":  patchErr,
 		})
 		return "", patchErr
 	}

 	patched, err := buf.Bytes()
 	if err != nil {
 		return "", err
 	}
 	if err := os.MkdirAll(filepath.Dir(input.Path), 0o700); err != nil {
 		return "", fmt.Errorf("failed to create directory %q: %w", filepath.Dir(input.Path), err)
 	}
 	if err := os.WriteFile(input.Path, patched, 0o600); err != nil {
 		return "", fmt.Errorf("failed to write patched contents to file %q: %w", input.Path, err)
 	}

 	response := new(strings.Builder)
 	fmt.Fprintf(response, "- Applied all patches\n")

 	if parsed {
 		parseErr := parseGo(patched)
 		if parseErr != nil {
 			return "", fmt.Errorf("after applying all patches, the file no longer parses:\n%w", parseErr)
 		}
 	}

 	if autogenerated {
 		fmt.Fprintf(response, "- WARNING: %q appears to be autogenerated. Patches were applied anyway.\n", input.Path)
 	}

 	// TODO: maybe report the patch result to the model, i.e. some/all of the new code after the patches and formatting.
 	return response.String(), nil
 }

 func parseGo(buf []byte) error {
 	fset := token.NewFileSet()
 	_, err := parser.ParseFile(fset, "", buf, parser.SkipObjectResolution)
 	return err
 }

 func isAutogeneratedGoFile(buf []byte) bool {
 	for _, sig := range autogeneratedSignals {
 		if bytes.Contains(buf, []byte(sig)) {
 			return true
 		}
 	}

 	// https://pkg.go.dev/cmd/go#hdr-Generate_Go_files_by_processing_source
 	// "This line must appear before the first non-comment, non-blank text in the file."
 	// Approximate that by looking for it at the top of the file, before the last of the imports.
 	// (Sometimes people put it after the package declaration, because of course they do.)
 	// At least in the imports region we know it's not part of their actual code;
 	// we don't want to ignore the generator (which also includes these strings!),
 	// just the generated code.
 	fset := token.NewFileSet()
 	f, err := parser.ParseFile(fset, "x.go", buf, parser.ImportsOnly|parser.ParseComments)
 	if err == nil {
 		for _, cg := range f.Comments {
 			t := strings.ToLower(cg.Text())
 			for _, sig := range autogeneratedHeaderSignals {
 				if strings.Contains(t, sig) {
 					return true
 				}
 			}
 		}
 	}

 	return false
 }

 // autogeneratedSignals are signals that a file is autogenerated, when present anywhere in the file.
 var autogeneratedSignals = [][]byte{
 	[]byte("\nfunc bindataRead("), // pre-embed bindata packed file
 }

 // autogeneratedHeaderSignals are signals that a file is autogenerated, when present at the top of the file.
 var autogeneratedHeaderSignals = []string{
 	// canonical would be `(?m)^// Code generated .* DO NOT EDIT\.$`
 	// but people screw it up, a lot, so be more lenient
 	strings.ToLower("generate"),
 	strings.ToLower("DO NOT EDIT"),
 	strings.ToLower("export by"),
 }
	package claudetool

	import (
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"go/parser"
	"go/token"
	"log/slog"
	"os"
	"path/filepath"
	"strings"

	"sketch.dev/ant"
	"sketch.dev/claudetool/editbuf"
	"sketch.dev/claudetool/patchkit"
	)

	// Patch is a tool for precise text modifications in files.
	var Patch = &ant.Tool{
	Name: PatchName,
	Description: strings.TrimSpace(PatchDescription),
	InputSchema: ant.MustSchema(PatchInputSchema),
	Run: PatchRun,
	}

	const (
	PatchName = "patch"
	PatchDescription = `
	File modification tool for precise text edits.

	Operations:
	- replace: Substitute text with new content
	- append_eof: Append new text at the end of the file
	- prepend_bof: Insert new text at the beginning of the file
	- overwrite: Replace the entire file with new content (automatically creates the file)

	Usage notes:
	- All inputs are interpreted literally (no automatic newline or whitespace handling)
	- For replace operations, oldText must appear EXACTLY ONCE in the file
	`

	// If you modify this, update the termui template for prettier rendering.
	PatchInputSchema = `
	{
	"type": "object",
	"required": ["path", "patches"],
	"properties": {
	"path": {
	"type": "string",
	"description": "Absolute path to the file to patch"
	},
	"patches": {
	"type": "array",
	"description": "List of patch requests to apply",
	"items": {
	"type": "object",
	"required": ["operation", "newText"],
	"properties": {
	"operation": {
	"type": "string",
	"enum": ["replace", "append_eof", "prepend_bof", "overwrite"],
	"description": "Type of operation to perform"
	},
	"oldText": {
	"type": "string",
	"description": "Text to locate for the operation (must be unique in file, required for replace)"
	},
	"newText": {
	"type": "string",
	"description": "The new text to use (empty for deletions)"
	}
	}
	}
	}
	}
	}
	`
	)

	// TODO: maybe rename PatchRequest to PatchOperation or PatchSpec or PatchPart or just Patch?

	type patchInput struct {
	Path string `json:"path"`
	Patches []patchRequest `json:"patches"`
	}

	type patchRequest struct {
	Operation string `json:"operation"`
	OldText string `json:"oldText,omitempty"`
	NewText string `json:"newText,omitempty"`
	}

	// PatchRun is the entry point for the user_patch tool.
	func PatchRun(ctx context.Context, m json.RawMessage) (string, error) {
	var input patchInput
	if err := json.Unmarshal(m, &input); err != nil {
	return "", fmt.Errorf("failed to unmarshal user_patch input: %w", err)
	}

	// Validate the input
	if !filepath.IsAbs(input.Path) {
	return "", fmt.Errorf("path %q is not absolute", input.Path)
	}
	if len(input.Patches) == 0 {
	return "", fmt.Errorf("no patches provided")
	}
	// TODO: check whether the file is autogenerated, and if so, require a "force" flag to modify it.

	orig, err := os.ReadFile(input.Path)
	// If the file doesn't exist, we can still apply patches
	// that don't require finding existing text.
	switch {
	case errors.Is(err, os.ErrNotExist):
	for _, patch := range input.Patches {
	switch patch.Operation {
	case "prepend_bof", "append_eof", "overwrite":
	default:
	return "", fmt.Errorf("file %q does not exist", input.Path)
	}
	}
	case err != nil:
	return "", fmt.Errorf("failed to read file %q: %w", input.Path, err)
	}

	likelyGoFile := strings.HasSuffix(input.Path, ".go")

	autogenerated := likelyGoFile && isAutogeneratedGoFile(orig)
	parsed := likelyGoFile && parseGo(orig) != nil

	origStr := string(orig)
	// Process the patches "simultaneously", minimizing them along the way.
	// Claude generates patches that interact with each other.
	buf := editbuf.NewBuffer(orig)

	// TODO: is it better to apply the patches that apply cleanly and report on the failures?
	// or instead have it be all-or-nothing?
	// For now, it is all-or-nothing.
	// TODO: when the model gets into a "cannot apply patch" cycle of doom, how do we get it unstuck?
	// Also: how do we detect that it's in a cycle?
	var patchErr error
	for i, patch := range input.Patches {
	switch patch.Operation {
	case "prepend_bof":
	buf.Insert(0, patch.NewText)
	case "append_eof":
	buf.Insert(len(orig), patch.NewText)
	case "overwrite":
	buf.Replace(0, len(orig), patch.NewText)
	case "replace":
	if patch.OldText == "" {
	return "", fmt.Errorf("patch %d: oldText cannot be empty for %s operation", i, patch.Operation)
	}

	// Attempt to apply the patch.
	spec, count := patchkit.Unique(origStr, patch.OldText, patch.NewText)
	switch count {
	case 0:
	// no matches, maybe recoverable, continued below
	case 1:
	// exact match, apply
	slog.DebugContext(ctx, "patch_applied", "method", "unique")
	spec.ApplyToEditBuf(buf)
	continue
	case 2:
	// multiple matches
	patchErr = errors.Join(patchErr, fmt.Errorf("old text not unique:\n%s", patch.OldText))
	default:
	// TODO: return an error instead of using agentPatch
	slog.ErrorContext(ctx, "unique returned unexpected count", "count", count)
	patchErr = errors.Join(patchErr, fmt.Errorf("internal error"))
	continue
	}

	// The following recovery mechanisms are heuristic.
	// They aren't perfect, but they appear safe,
	// and the cases they cover appear with some regularity.

	// Try adjusting the whitespace prefix.
	spec, ok := patchkit.UniqueDedent(origStr, patch.OldText, patch.NewText)
	if ok {
	slog.DebugContext(ctx, "patch_applied", "method", "unique_dedent")
	spec.ApplyToEditBuf(buf)
	continue
	}

	// Try ignoring leading/trailing whitespace in a semantically safe way.
	spec, ok = patchkit.UniqueInValidGo(origStr, patch.OldText, patch.NewText)
	if ok {
	slog.DebugContext(ctx, "patch_applied", "method", "unique_in_valid_go")
	spec.ApplyToEditBuf(buf)
	continue
	}

	// Try ignoring semantically insignificant whitespace.
	spec, ok = patchkit.UniqueGoTokens(origStr, patch.OldText, patch.NewText)
	if ok {
	slog.DebugContext(ctx, "patch_applied", "method", "unique_go_tokens")
	spec.ApplyToEditBuf(buf)
	continue
	}

	// Try trimming the first line of the patch, if we can do so safely.
	spec, ok = patchkit.UniqueTrim(origStr, patch.OldText, patch.NewText)
	if ok {
	slog.DebugContext(ctx, "patch_applied", "method", "unique_trim")
	spec.ApplyToEditBuf(buf)
	continue
	}

	// No dice.
	patchErr = errors.Join(patchErr, fmt.Errorf("old text not found:\n%s", patch.OldText))
	continue
	default:
	return "", fmt.Errorf("unrecognized operation %q", patch.Operation)
	}
	}

	if patchErr != nil {
	sendTelemetry(ctx, "patch_error", map[string]any{
	"orig": origStr,
	"patches": input.Patches,
	"errors": patchErr,
	})
	return "", patchErr
	}

	patched, err := buf.Bytes()
	if err != nil {
	return "", err
	}
	if err := os.MkdirAll(filepath.Dir(input.Path), 0o700); err != nil {
	return "", fmt.Errorf("failed to create directory %q: %w", filepath.Dir(input.Path), err)
	}
	if err := os.WriteFile(input.Path, patched, 0o600); err != nil {
	return "", fmt.Errorf("failed to write patched contents to file %q: %w", input.Path, err)
	}

	response := new(strings.Builder)
	fmt.Fprintf(response, "- Applied all patches\n")

	if parsed {
	parseErr := parseGo(patched)
	if parseErr != nil {
	return "", fmt.Errorf("after applying all patches, the file no longer parses:\n%w", parseErr)
	}
	}

	if autogenerated {
	fmt.Fprintf(response, "- WARNING: %q appears to be autogenerated. Patches were applied anyway.\n", input.Path)
	}

	// TODO: maybe report the patch result to the model, i.e. some/all of the new code after the patches and formatting.
	return response.String(), nil
	}

	func parseGo(buf []byte) error {
	fset := token.NewFileSet()
	_, err := parser.ParseFile(fset, "", buf, parser.SkipObjectResolution)
	return err
	}

	func isAutogeneratedGoFile(buf []byte) bool {
	for _, sig := range autogeneratedSignals {
	if bytes.Contains(buf, []byte(sig)) {
	return true
	}
	}

	// https://pkg.go.dev/cmd/go#hdr-Generate_Go_files_by_processing_source
	// "This line must appear before the first non-comment, non-blank text in the file."
	// Approximate that by looking for it at the top of the file, before the last of the imports.
	// (Sometimes people put it after the package declaration, because of course they do.)
	// At least in the imports region we know it's not part of their actual code;
	// we don't want to ignore the generator (which also includes these strings!),
	// just the generated code.
	fset := token.NewFileSet()
	f, err := parser.ParseFile(fset, "x.go", buf, parser.ImportsOnly\|parser.ParseComments)
	if err == nil {
	for _, cg := range f.Comments {
	t := strings.ToLower(cg.Text())
	for _, sig := range autogeneratedHeaderSignals {
	if strings.Contains(t, sig) {
	return true
	}
	}
	}
	}

	return false
	}

	// autogeneratedSignals are signals that a file is autogenerated, when present anywhere in the file.
	var autogeneratedSignals = [][]byte{
	[]byte("\nfunc bindataRead("), // pre-embed bindata packed file
	}

	// autogeneratedHeaderSignals are signals that a file is autogenerated, when present at the top of the file.
	var autogeneratedHeaderSignals = []string{
	// canonical would be `(?m)^// Code generated .* DO NOT EDIT\.$`
	// but people screw it up, a lot, so be more lenient
	strings.ToLower("generate"),
	strings.ToLower("DO NOT EDIT"),
	strings.ToLower("export by"),
	}