154 lines
4.3 KiB
Go
154 lines
4.3 KiB
Go
package tools
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"errors"
|
|
"fmt"
|
|
"strings"
|
|
|
|
"agentu/pkg/llm"
|
|
)
|
|
|
|
// ────────────────────────────────────────────────────────────────────────────
|
|
// ApplyPatchTool — multi-file patch using V4 format
|
|
// ────────────────────────────────────────────────────────────────────────────
|
|
|
|
type ApplyPatchTool struct {
|
|
workingDir string
|
|
}
|
|
|
|
func NewApplyPatchTool(workingDir string) *ApplyPatchTool {
|
|
return &ApplyPatchTool{workingDir: workingDir}
|
|
}
|
|
|
|
// MutatingTool interface.
|
|
|
|
func (*ApplyPatchTool) Mutates() bool { return true }
|
|
|
|
// CallMutatingTool — dry_run is non-mutating.
|
|
func (t *ApplyPatchTool) MutatesCall(args json.RawMessage) bool {
|
|
var a applyPatchArgs
|
|
if err := json.Unmarshal(args, &a); err != nil {
|
|
return true // conservative
|
|
}
|
|
return !a.DryRun
|
|
}
|
|
|
|
// PreviewableTool interface.
|
|
|
|
func (t *ApplyPatchTool) Preview(ctx context.Context, args json.RawMessage) (string, error) {
|
|
var a applyPatchArgs
|
|
if err := json.Unmarshal(args, &a); err != nil {
|
|
return "", fmt.Errorf("parse arguments: %w", err)
|
|
}
|
|
return summarizePatch(a.Patch)
|
|
}
|
|
|
|
// TargetPathTool interface.
|
|
|
|
func (t *ApplyPatchTool) TargetPaths(args json.RawMessage) ([]string, error) {
|
|
var a applyPatchArgs
|
|
if err := json.Unmarshal(args, &a); err != nil {
|
|
return nil, fmt.Errorf("parse arguments: %w", err)
|
|
}
|
|
return extractTargetPaths(t.workingDir, a.Patch)
|
|
}
|
|
|
|
// Definition.
|
|
|
|
func (t *ApplyPatchTool) Definition() llm.Tool {
|
|
return llm.Tool{
|
|
Type: "function",
|
|
Function: llm.ToolFunction{
|
|
Name: "apply_patch",
|
|
Description: "Apply a multi-file patch to create, update, or delete files. Uses a V4 diff format with *** Begin Patch / *** Add File / *** Update File / *** Delete File markers. Prefer this over multiple file_edit calls when changing 2+ files.",
|
|
Parameters: JSONSchema(`{
|
|
"type": "object",
|
|
"properties": {
|
|
"patch": {"type": "string", "description": "V4 format patch text. Must start with *** Begin Patch and end with *** End Patch."},
|
|
"dry_run": {"type": "boolean", "description": "Preview the patch without writing any files."}
|
|
},
|
|
"required": ["patch"],
|
|
"additionalProperties": false
|
|
}`),
|
|
},
|
|
}
|
|
}
|
|
|
|
type applyPatchArgs struct {
|
|
Patch string `json:"patch"`
|
|
DryRun bool `json:"dry_run"`
|
|
}
|
|
|
|
// Execute.
|
|
|
|
func (t *ApplyPatchTool) Execute(ctx context.Context, raw json.RawMessage) (string, error) {
|
|
var args applyPatchArgs
|
|
if err := json.Unmarshal(raw, &args); err != nil {
|
|
return "", fmt.Errorf("parse arguments: %w", err)
|
|
}
|
|
|
|
if strings.TrimSpace(args.Patch) == "" {
|
|
return "", fmt.Errorf("patch is empty")
|
|
}
|
|
|
|
pp, err := parsePatch(args.Patch)
|
|
if err != nil {
|
|
return "", err
|
|
}
|
|
|
|
if args.DryRun {
|
|
summary, err := summarizePatch(args.Patch)
|
|
if err != nil {
|
|
return "", err
|
|
}
|
|
return "dry run — " + summary, nil
|
|
}
|
|
|
|
rollback := NewRollbackStore()
|
|
var snapshotIDs []string
|
|
for _, pf := range pp.files {
|
|
path, err := resolvePathForWrite(t.workingDir, pf.path)
|
|
if err != nil {
|
|
return "", err
|
|
}
|
|
id, err := rollback.Capture("apply_patch", pf.path, path)
|
|
if err != nil {
|
|
return "", err
|
|
}
|
|
snapshotIDs = append(snapshotIDs, id)
|
|
}
|
|
|
|
var results []string
|
|
for _, pf := range pp.files {
|
|
if err := ctx.Err(); err != nil {
|
|
if rollbackErr := restoreSnapshots(rollback, snapshotIDs); rollbackErr != nil {
|
|
return "", fmt.Errorf("%w; rollback failed: %v", err, rollbackErr)
|
|
}
|
|
return "", err
|
|
}
|
|
result, err := applyPatchFile(t.workingDir, pf)
|
|
if err != nil {
|
|
if rollbackErr := restoreSnapshots(rollback, snapshotIDs); rollbackErr != nil {
|
|
return "", fmt.Errorf("apply patch: %w; rollback failed: %v", err, rollbackErr)
|
|
}
|
|
return "", fmt.Errorf("apply patch: %w", err)
|
|
}
|
|
results = append(results, result)
|
|
}
|
|
|
|
summary := fmt.Sprintf("applied patch: %d file(s) changed\n%s", len(results), strings.Join(results, "\n"))
|
|
return summary, nil
|
|
}
|
|
|
|
func restoreSnapshots(store *RollbackStore, ids []string) error {
|
|
var errs []error
|
|
for i := len(ids) - 1; i >= 0; i-- {
|
|
if _, err := store.Restore(ids[i]); err != nil {
|
|
errs = append(errs, err)
|
|
}
|
|
}
|
|
return errors.Join(errs...)
|
|
}
|