// Package edit provides AST-aware file editing capabilities. package edit import ( "bytes" "context" "fmt" "os" "regexp" "strings" "sync" "github.com/lukaszraczylo/mcp-filepuff/internal/parser" "github.com/lukaszraczylo/mcp-filepuff/pkg/errors" "github.com/lukaszraczylo/mcp-filepuff/pkg/protocol" "github.com/sergi/go-diff/diffmatchpatch" sitter "github.com/smacker/go-tree-sitter" ) // Global regex cache for compiled patterns (thread-safe) var regexCache sync.Map // string -> *regexp.Regexp // compileRegex compiles a regex pattern with caching for performance. func compileRegex(pattern string) (*regexp.Regexp, error) { // Check cache first if cached, ok := regexCache.Load(pattern); ok { return cached.(*regexp.Regexp), nil } // Compile and cache re, err := regexp.Compile(pattern) if err != nil { return nil, err } regexCache.Store(pattern, re) return re, nil } // EditOperation defines the type of edit operation. type EditOperation string const ( EditReplace EditOperation = "replace" EditInsertBefore EditOperation = "insert_before" EditInsertAfter EditOperation = "insert_after" EditDelete EditOperation = "delete" ) // ASTEdit represents an AST-aware edit request. type ASTEdit struct { File string `json:"file"` Operation EditOperation `json:"operation"` NewContent string `json:"new_content,omitempty"` Selector ASTSelector `json:"selector"` } // ASTSelector specifies how to find the target node. type ASTSelector struct { Kind string `json:"kind,omitempty"` Name string `json:"name,omitempty"` Pattern string `json:"pattern,omitempty"` Text string `json:"text,omitempty"` TextPattern string `json:"text_pattern,omitempty"` AtLine int `json:"at_line,omitempty"` Index int `json:"index,omitempty"` LineEnd int `json:"line_end,omitempty"` } // EditResult contains the result of an edit operation. type EditResult struct { Diff string `json:"diff,omitempty"` OriginalContent string `json:"original_content,omitempty"` NewContent string `json:"new_content,omitempty"` Error string `json:"error,omitempty"` Success bool `json:"success"` Applied bool `json:"applied"` } // Engine performs AST-aware edits. type Engine struct { registry *parser.Registry fileLocks sync.Map // map[string]*sync.Mutex for per-file locking } // NewEngine creates a new edit engine. func NewEngine(registry *parser.Registry) *Engine { return &Engine{ registry: registry, fileLocks: sync.Map{}, } } // lockFile acquires a lock for the specified file and returns an unlock function. // This prevents concurrent edits to the same file which could cause corruption. func (e *Engine) lockFile(filePath string) func() { // Get or create mutex for this file actual, _ := e.fileLocks.LoadOrStore(filePath, &sync.Mutex{}) mu := actual.(*sync.Mutex) mu.Lock() return mu.Unlock } // Preview generates a preview of an edit without applying it. func (e *Engine) Preview(ctx context.Context, edit *ASTEdit) (*EditResult, error) { return e.performEdit(ctx, edit, false) } // Apply performs an edit and writes the result to disk. // Uses file locking to prevent concurrent edits to the same file. func (e *Engine) Apply(ctx context.Context, edit *ASTEdit) (*EditResult, error) { unlock := e.lockFile(edit.File) defer unlock() return e.performEdit(ctx, edit, true) } // performEdit executes an edit operation. func (e *Engine) performEdit(ctx context.Context, edit *ASTEdit, apply bool) (*EditResult, error) { // Determine if we should use text mode useTextMode := e.shouldUseTextMode(edit) if useTextMode { return e.performTextEdit(ctx, edit, apply) } return e.performASTEdit(ctx, edit, apply) } // shouldUseTextMode determines if text-based editing should be used. func (e *Engine) shouldUseTextMode(edit *ASTEdit) bool { // Use text mode if text-specific selectors are provided if edit.Selector.Text != "" || edit.Selector.TextPattern != "" { return true } // Use text mode if line range is specified without AST selectors if edit.Selector.AtLine > 0 && edit.Selector.LineEnd > 0 && edit.Selector.Kind == "" && edit.Selector.Name == "" && edit.Selector.Pattern == "" { return true } // Use text mode if language is not supported for AST lang := protocol.DetectLanguage(edit.File) return lang == protocol.LangUnknown } // performASTEdit executes an AST-aware edit operation. func (e *Engine) performASTEdit(ctx context.Context, edit *ASTEdit, apply bool) (*EditResult, error) { // Validate operation if err := e.validateASTEdit(edit); err != nil { return &EditResult{Success: false, Error: err.Error()}, nil } // Read file content, err := os.ReadFile(edit.File) if err != nil { structuredErr := errors.NewFileNotReadableError(edit.File, err) return &EditResult{Success: false, Error: structuredErr.Error()}, nil } // Parse file parseResult, err := e.registry.Parse(ctx, edit.File, content) if err != nil { return &EditResult{Success: false, Error: err.Error()}, nil } // Find target node node, err := e.resolveSelector(edit.Selector, parseResult.Tree, content) if err != nil { return &EditResult{Success: false, Error: err.Error()}, nil } // Apply edit newContent, err := e.applyEdit(edit, node, content) if err != nil { return &EditResult{Success: false, Error: err.Error()}, nil } // Validate new content (re-parse) _, err = e.registry.Parse(ctx, edit.File, newContent) if err != nil { structuredErr := errors.NewEditValidationError(edit.File, err) return &EditResult{ Success: false, Error: structuredErr.Error(), }, nil } // Generate diff diff := generateDiff(string(content), string(newContent), edit.File) result := &EditResult{ Success: true, Diff: diff, OriginalContent: string(content), NewContent: string(newContent), Applied: false, } // Apply changes if requested if apply { if err := os.WriteFile(edit.File, newContent, 0600); err != nil { structuredErr := errors.NewFileNotWritableError(edit.File, err) return &EditResult{ Success: false, Error: structuredErr.Error(), }, nil } result.Applied = true } return result, nil } // performTextEdit executes a text-based edit operation for non-AST files. func (e *Engine) performTextEdit(_ context.Context, edit *ASTEdit, apply bool) (*EditResult, error) { // Validate operation if err := e.validateTextEdit(edit); err != nil { return &EditResult{Success: false, Error: err.Error()}, nil } // Read file content, err := os.ReadFile(edit.File) if err != nil { structuredErr := errors.NewFileNotReadableError(edit.File, err) return &EditResult{Success: false, Error: structuredErr.Error()}, nil } // Find the text selection (byte range) start, end, err := e.resolveTextSelector(edit.Selector, content) if err != nil { return &EditResult{Success: false, Error: err.Error()}, nil } // Apply edit newContent, err := e.applyTextEditOperation(edit.Operation, content, start, end, edit.NewContent) if err != nil { return &EditResult{Success: false, Error: err.Error()}, nil } // Generate diff diff := generateDiff(string(content), string(newContent), edit.File) result := &EditResult{ Success: true, Diff: diff, OriginalContent: string(content), NewContent: string(newContent), Applied: false, } // Apply changes if requested if apply { if err := os.WriteFile(edit.File, newContent, 0600); err != nil { structuredErr := errors.NewFileNotWritableError(edit.File, err) return &EditResult{ Success: false, Error: structuredErr.Error(), }, nil } result.Applied = true } return result, nil } // validateBaseEdit checks common edit request fields. func (e *Engine) validateBaseEdit(edit *ASTEdit) error { if edit.File == "" { return errors.NewInvalidEditError("file is required") } if edit.Operation == "" { return errors.NewInvalidEditError("operation is required") } // Validate operation type switch edit.Operation { case EditReplace, EditInsertBefore, EditInsertAfter: if edit.NewContent == "" { return errors.NewInvalidEditError(fmt.Sprintf("new_content is required for %s operation", edit.Operation)) } case EditDelete: // new_content not required default: return errors.NewInvalidEditError(fmt.Sprintf("unknown operation: %s", edit.Operation)) } return nil } // validateASTEdit checks if an AST edit request is valid. func (e *Engine) validateASTEdit(edit *ASTEdit) error { if err := e.validateBaseEdit(edit); err != nil { return err } // Validate AST selector if edit.Selector.Kind == "" && edit.Selector.Name == "" && edit.Selector.Pattern == "" && edit.Selector.AtLine == 0 { return errors.NewInvalidEditError("AST selector must specify at least one of: kind, name, pattern, or at_line") } return nil } // validateTextEdit checks if a text edit request is valid. func (e *Engine) validateTextEdit(edit *ASTEdit) error { if err := e.validateBaseEdit(edit); err != nil { return err } // Validate text selector - need at least one text selection method hasTextSelector := edit.Selector.Text != "" || edit.Selector.TextPattern != "" || edit.Selector.AtLine > 0 if !hasTextSelector { return errors.NewInvalidEditError("text selector must specify at least one of: text, text_pattern, or at_line") } // Validate regex pattern if provided (uses cached compilation) if edit.Selector.TextPattern != "" { if _, err := compileRegex(edit.Selector.TextPattern); err != nil { return errors.Wrap(errors.ErrInvalidEdit, "invalid text_pattern regex", err) } } return nil } // resolveSelector finds the target node based on the selector. func (e *Engine) resolveSelector(sel ASTSelector, tree *sitter.Tree, content []byte) (*sitter.Node, error) { if tree == nil { return nil, errors.NewNodeNotFoundError("no AST tree available") } root := tree.RootNode() if root == nil { return nil, errors.NewNodeNotFoundError("empty AST tree") } var matches []*sitter.Node parser.WalkTree(root, func(n *sitter.Node) bool { if e.matchesSelector(sel, n, content) { matches = append(matches, n) } return true }) if len(matches) == 0 { selectorDesc := fmt.Sprintf("kind=%s name=%s pattern=%s line=%d", sel.Kind, sel.Name, sel.Pattern, sel.AtLine) return nil, errors.NewNodeNotFoundError(selectorDesc) } // Use index to select specific match index := sel.Index if index < 0 || index >= len(matches) { return nil, errors.NewInvalidSelectionError(fmt.Sprintf("selector matched %d nodes, but index %d is out of range", len(matches), index)) } return matches[index], nil } // matchesSelector checks if a node matches the selector criteria. func (e *Engine) matchesSelector(sel ASTSelector, n *sitter.Node, content []byte) bool { // Check kind if sel.Kind != "" && n.Type() != sel.Kind { return false } // Check name (look for identifier in the node) if sel.Name != "" { nameNode := n.ChildByFieldName("name") if nameNode == nil { // Also try to find an identifier child found := false for i := 0; i < int(n.NamedChildCount()); i++ { child := n.NamedChild(i) if child != nil && child.Type() == "identifier" { if parser.GetNodeText(child, content) == sel.Name { found = true break } } } if !found { return false } } else if parser.GetNodeText(nameNode, content) != sel.Name { return false } } // Check line if sel.AtLine > 0 { startLine := int(n.StartPoint().Row) + 1 endLine := int(n.EndPoint().Row) + 1 if sel.AtLine < startLine || sel.AtLine > endLine { return false } } // Pattern matching is handled separately (simplified here) if sel.Pattern != "" { nodeText := parser.GetNodeText(n, content) if !strings.Contains(nodeText, sel.Pattern) { return false } } return true } // applyEdit applies the edit operation to the content. func (e *Engine) applyEdit(edit *ASTEdit, node *sitter.Node, content []byte) ([]byte, error) { startByte := node.StartByte() endByte := node.EndByte() // Detect and preserve indentation indentation := detectIndentation(content, startByte) newContent := indentContent(edit.NewContent, indentation) var result []byte switch edit.Operation { case EditReplace: result = append(result, content[:startByte]...) result = append(result, []byte(newContent)...) result = append(result, content[endByte:]...) case EditInsertBefore: result = append(result, content[:startByte]...) result = append(result, []byte(newContent)...) result = append(result, '\n') result = append(result, content[startByte:]...) case EditInsertAfter: result = append(result, content[:endByte]...) result = append(result, '\n') result = append(result, []byte(newContent)...) result = append(result, content[endByte:]...) case EditDelete: result = append(result, content[:startByte]...) result = append(result, content[endByte:]...) default: return nil, errors.NewInvalidEditError(fmt.Sprintf("unknown operation: %s", edit.Operation)) } return result, nil } // detectIndentation detects the indentation at a given byte position. func detectIndentation(content []byte, bytePos uint32) string { // Find the start of the line lineStart := int(bytePos) for lineStart > 0 && content[lineStart-1] != '\n' { lineStart-- } // Extract leading whitespace var indent strings.Builder for i := lineStart; i < int(bytePos) && i < len(content); i++ { c := content[i] if c == ' ' || c == '\t' { indent.WriteByte(c) } else { break } } return indent.String() } // indentContent applies indentation to multi-line content. func indentContent(content string, indent string) string { if indent == "" { return content } lines := strings.Split(content, "\n") for i, line := range lines { if i > 0 && line != "" { lines[i] = indent + line } } return strings.Join(lines, "\n") } // generateDiff creates a unified diff between original and modified content. // Uses Myers diff algorithm for accurate and readable diffs. func generateDiff(original, modified, filename string) string { dmp := diffmatchpatch.New() diffs := dmp.DiffMain(original, modified, false) // Cleanup for readability diffs = dmp.DiffCleanupSemantic(diffs) // Convert to unified diff format var buf bytes.Buffer buf.WriteString(fmt.Sprintf("--- %s\n", filename)) buf.WriteString(fmt.Sprintf("+++ %s\n", filename)) // Group diffs into hunks lineNum := 1 for _, diff := range diffs { lines := strings.Split(diff.Text, "\n") for i, line := range lines { // Skip empty last line from split if i == len(lines)-1 && line == "" { continue } switch diff.Type { case diffmatchpatch.DiffDelete: buf.WriteString(fmt.Sprintf("-%s\n", line)) case diffmatchpatch.DiffInsert: buf.WriteString(fmt.Sprintf("+%s\n", line)) case diffmatchpatch.DiffEqual: buf.WriteString(fmt.Sprintf(" %s\n", line)) lineNum++ } } } return buf.String() } // resolveTextSelector finds the byte range for a text-based selection. func (e *Engine) resolveTextSelector(sel ASTSelector, content []byte) (start, end int, err error) { switch { case sel.Text != "": return e.findExactText(content, sel.Text, sel.Index) case sel.TextPattern != "": return e.findRegexPattern(content, sel.TextPattern, sel.Index) case sel.AtLine > 0: return e.findLineRange(content, sel.AtLine, sel.LineEnd) default: return 0, 0, errors.NewInvalidEditError("text selector requires text, text_pattern, or at_line") } } // findExactText finds an exact text match in content. func (e *Engine) findExactText(content []byte, text string, index int) (start, end int, err error) { if text == "" { return 0, 0, errors.NewInvalidEditError("text selector cannot be empty") } textBytes := []byte(text) type match struct{ start, end int } var matches []match offset := 0 for { idx := bytes.Index(content[offset:], textBytes) if idx == -1 { break } matches = append(matches, match{ start: offset + idx, end: offset + idx + len(textBytes), }) offset += idx + 1 } if len(matches) == 0 { return 0, 0, errors.NewInvalidSelectionError(fmt.Sprintf("text not found: %q", truncateString(text, 50))) } // If multiple matches and no index specified, require explicit selection if len(matches) > 1 && index == 0 { // Check if index was explicitly set to 0 or just defaulted // Since we can't distinguish, we'll allow index 0 but warn about multiple matches // Actually, let's be strict and require explicit index for multiple matches locations := make([]string, 0, min(len(matches), 5)) for i, m := range matches { if i >= 5 { locations = append(locations, fmt.Sprintf("... and %d more", len(matches)-5)) break } line := countLines(content[:m.start]) + 1 locations = append(locations, fmt.Sprintf("line %d", line)) } return 0, 0, errors.NewInvalidSelectionError(fmt.Sprintf("text matches %d locations (%s); use selector_index to specify which one (0-%d)", len(matches), strings.Join(locations, ", "), len(matches)-1)) } if index >= len(matches) { return 0, 0, errors.NewInvalidSelectionError(fmt.Sprintf("selector_index %d out of range (found %d matches)", index, len(matches))) } return matches[index].start, matches[index].end, nil } // findRegexPattern finds a regex pattern match in content. func (e *Engine) findRegexPattern(content []byte, pattern string, index int) (start, end int, err error) { re, err := compileRegex(pattern) if err != nil { return 0, 0, errors.Wrap(errors.ErrInvalidEdit, "invalid regex pattern", err) } matches := re.FindAllIndex(content, -1) if len(matches) == 0 { return 0, 0, errors.NewInvalidSelectionError(fmt.Sprintf("pattern not found: %q", truncateString(pattern, 50))) } // If multiple matches and index is 0 (default), show error with locations if len(matches) > 1 && index == 0 { locations := make([]string, 0, min(len(matches), 5)) for i, m := range matches { if i >= 5 { locations = append(locations, fmt.Sprintf("... and %d more", len(matches)-5)) break } line := countLines(content[:m[0]]) + 1 locations = append(locations, fmt.Sprintf("line %d", line)) } return 0, 0, errors.NewInvalidSelectionError(fmt.Sprintf("pattern matches %d locations (%s); use selector_index to specify which one (0-%d)", len(matches), strings.Join(locations, ", "), len(matches)-1)) } if index >= len(matches) { return 0, 0, errors.NewInvalidSelectionError(fmt.Sprintf("selector_index %d out of range (found %d matches)", index, len(matches))) } return matches[index][0], matches[index][1], nil } // findLineRange finds the byte range for a line range selection. func (e *Engine) findLineRange(content []byte, lineStart, lineEnd int) (start, end int, err error) { if lineEnd == 0 { lineEnd = lineStart } if lineStart < 1 { return 0, 0, errors.NewInvalidEditError(fmt.Sprintf("line number must be >= 1, got %d", lineStart)) } if lineEnd < lineStart { return 0, 0, errors.NewInvalidEditError(fmt.Sprintf("line_end (%d) must be >= line (%d)", lineEnd, lineStart)) } lines := bytes.Split(content, []byte("\n")) totalLines := len(lines) // Convert to 0-indexed startIdx := lineStart - 1 endIdx := lineEnd - 1 if startIdx >= totalLines { return 0, 0, errors.NewInvalidSelectionError(fmt.Sprintf("line %d out of range (file has %d lines)", lineStart, totalLines)) } if endIdx >= totalLines { return 0, 0, errors.NewInvalidSelectionError(fmt.Sprintf("line_end %d out of range (file has %d lines)", lineEnd, totalLines)) } // Calculate byte positions start = 0 for i := 0; i < startIdx; i++ { start += len(lines[i]) + 1 // +1 for newline } end = start for i := startIdx; i <= endIdx; i++ { end += len(lines[i]) if i < totalLines-1 { end += 1 // newline } } return start, end, nil } // applyTextEditOperation applies a text edit operation. func (e *Engine) applyTextEditOperation(op EditOperation, content []byte, start, end int, newContent string) ([]byte, error) { // Detect indentation at the selection point indentation := detectIndentationAtByte(content, start) indentedContent := indentContent(newContent, indentation) var result []byte switch op { case EditReplace: result = append(result, content[:start]...) result = append(result, []byte(indentedContent)...) result = append(result, content[end:]...) case EditInsertBefore: result = append(result, content[:start]...) result = append(result, []byte(indentedContent)...) result = append(result, '\n') result = append(result, content[start:]...) case EditInsertAfter: result = append(result, content[:end]...) result = append(result, '\n') result = append(result, []byte(indentedContent)...) result = append(result, content[end:]...) case EditDelete: result = append(result, content[:start]...) result = append(result, content[end:]...) default: return nil, errors.NewInvalidEditError(fmt.Sprintf("unknown operation: %s", op)) } return result, nil } // detectIndentationAtByte detects indentation at a byte position. func detectIndentationAtByte(content []byte, bytePos int) string { // Find the start of the line lineStart := bytePos for lineStart > 0 && content[lineStart-1] != '\n' { lineStart-- } // Extract leading whitespace var indent strings.Builder for i := lineStart; i < bytePos && i < len(content); i++ { c := content[i] if c == ' ' || c == '\t' { indent.WriteByte(c) } else { break } } return indent.String() } // truncateString truncates a string to maxLen with ellipsis. func truncateString(s string, maxLen int) string { if len(s) <= maxLen { return s } return s[:maxLen-3] + "..." } // countLines counts the number of newlines in content. func countLines(content []byte) int { return bytes.Count(content, []byte("\n")) } // ValidateLanguage checks if AST editing is supported for a file. // Returns nil for supported languages, error for unsupported. // Note: Text-based editing is always available regardless of this check. func ValidateLanguage(filename string) error { lang := protocol.DetectLanguage(filename) if lang == protocol.LangUnknown { return fmt.Errorf("unsupported file type for AST editing: %s (text-based editing is available)", filename) } return nil }