feat: complete v0.3.1 agent workflow

internal/tools/rollback.go

@@ -0,0 +1,117 @@
+package tools
+
+import (
+	"crypto/rand"
+	"encoding/hex"
+	"fmt"
+	"io/fs"
+	"os"
+	"sync"
+	"time"
+
+	"agentu/internal/safeio"
+)
+
+// RollbackSnapshot captures the state of a file before a mutation.
+type RollbackSnapshot struct {
+	ID           string
+	Tool         string
+	Path         string
+	ResolvedPath string
+	Existed      bool
+	Data         []byte
+	Mode         fs.FileMode
+	CreatedAt    time.Time
+}
+
+// RollbackStore manages in-memory rollback snapshots for the current process.
+type RollbackStore struct {
+	mu        sync.Mutex
+	snapshots []RollbackSnapshot
+}
+
+// NewRollbackStore creates an empty rollback store.
+func NewRollbackStore() *RollbackStore {
+	return &RollbackStore{}
+}
+
+// generateSnapshotID creates a rollback ID: "rb_" + 12 lowercase hex chars.
+func generateSnapshotID() (string, error) {
+	b := make([]byte, 6)
+	if _, err := rand.Read(b); err != nil {
+		return "", fmt.Errorf("generate snapshot id: %w", err)
+	}
+	return "rb_" + hex.EncodeToString(b), nil
+}
+
+// Capture stores the current state of the file at resolvedPath before a tool
+// modifies it. Returns the snapshot ID.
+func (s *RollbackStore) Capture(tool string, displayPath string, resolvedPath string) (string, error) {
+	id, err := generateSnapshotID()
+	if err != nil {
+		return "", err
+	}
+
+	snap := RollbackSnapshot{
+		ID:           id,
+		Tool:         tool,
+		Path:         displayPath,
+		ResolvedPath: resolvedPath,
+		CreatedAt:    time.Now().UTC(),
+	}
+
+	info, err := os.Stat(resolvedPath)
+	if err != nil {
+		if os.IsNotExist(err) {
+			snap.Existed = false
+		} else {
+			return "", fmt.Errorf("stat file: %w", err)
+		}
+	} else {
+		snap.Existed = true
+		snap.Mode = info.Mode().Perm()
+		data, err := os.ReadFile(resolvedPath)
+		if err != nil {
+			return "", fmt.Errorf("read file for snapshot: %w", err)
+		}
+		snap.Data = data
+	}
+
+	s.mu.Lock()
+	s.snapshots = append(s.snapshots, snap)
+	s.mu.Unlock()
+	return id, nil
+}
+
+// Restore rolls back a file to the state captured in the given snapshot.
+func (s *RollbackStore) Restore(id string) (RollbackSnapshot, error) {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	for i := len(s.snapshots) - 1; i >= 0; i-- {
+		if s.snapshots[i].ID == id {
+			snap := s.snapshots[i]
+			if snap.Existed {
+				if err := safeio.AtomicWriteFile(snap.ResolvedPath, snap.Data, snap.Mode); err != nil {
+					return RollbackSnapshot{}, fmt.Errorf("restore file: %w", err)
+				}
+			} else {
+				if err := os.Remove(snap.ResolvedPath); err != nil && !os.IsNotExist(err) {
+					return RollbackSnapshot{}, fmt.Errorf("remove created file: %w", err)
+				}
+			}
+			s.snapshots = append(s.snapshots[:i], s.snapshots[i+1:]...)
+			return snap, nil
+		}
+	}
+	return RollbackSnapshot{}, fmt.Errorf("rollback snapshot %s not found", id)
+}
+
+// LastID returns the most recent snapshot ID, or "" if empty.
+func (s *RollbackStore) LastID() string {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	if len(s.snapshots) == 0 {
+		return ""
+	}
+	return s.snapshots[len(s.snapshots)-1].ID
+}