118 lines
2.9 KiB
Go
118 lines
2.9 KiB
Go
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
|
|
}
|