agent-Specialization/modules/shallow_versioning.py

"""Conversation-scoped shallow file versioning for workspace edits.

Mirrors Claude Code's file-history approach:
- Only files touched by write_file / edit_file are tracked.
- Before each edit the current content is backed up.
- A snapshot is taken per user message, referencing the latest backup version
  of every tracked file at that point in time.
- Rewind restores the tracked files to the state recorded for a target message.
"""

from __future__ import annotations

import hashlib
import json
import shutil
from dataclasses import dataclass, field
from datetime import datetime
from pathlib import Path
from typing import Any, Dict, List, Optional, Set


class ShallowVersioningError(RuntimeError):
    """Raised when shallow versioning fails."""


@dataclass
class FileBackup:
    backup_file_name: Optional[str]  # null means the file did not exist at that version
    version: int
    backup_time: str


@dataclass
class ShallowSnapshot:
    message_id: str
    tracked_file_backups: Dict[str, FileBackup]
    timestamp: str


class ShallowVersioningManager:
    """Lightweight per-conversation backup for files edited by AI tools."""

    MAX_SNAPSHOTS = 100

    def __init__(self, project_path: Path | str, data_dir: Path | str, conversation_id: str):
        self.project_path = Path(project_path).expanduser().resolve()
        self.conversation_id = str(conversation_id or "").strip()
        if not self.conversation_id:
            raise ShallowVersioningError("缺少 conversation_id")
        self.save_root = (Path(data_dir).expanduser().resolve() / "save" / self.conversation_id).resolve()
        self.backup_dir = self.save_root / "shallow_backups"
        self.state_file = self.backup_dir / "state.jsonl"
        self._tracked_files: Set[str] = set()
        self._snapshots: List[ShallowSnapshot] = []
        self._snapshot_sequence = 0
        self._load_state()

    # ----------------------------
    # persistence
    # ----------------------------
    def _load_state(self) -> None:
        if not self.state_file.exists():
            return
        try:
            for raw in self.state_file.read_text(encoding="utf-8", errors="ignore").splitlines():
                line = raw.strip()
                if not line:
                    continue
                try:
                    row = json.loads(line)
                except Exception:
                    continue
                if not isinstance(row, dict):
                    continue
                snapshot = self._deserialize_snapshot(row)
                if snapshot:
                    self._snapshots.append(snapshot)
                    self._snapshot_sequence += 1
                    for path in snapshot.tracked_file_backups:
                        self._tracked_files.add(path)
        except OSError:
            pass

    def _persist_snapshot(self, snapshot: ShallowSnapshot) -> None:
        self.backup_dir.mkdir(parents=True, exist_ok=True)
        with self.state_file.open("a", encoding="utf-8") as fh:
            fh.write(json.dumps(self._serialize_snapshot(snapshot), ensure_ascii=False) + "\n")

    def _serialize_snapshot(self, snapshot: ShallowSnapshot) -> Dict[str, Any]:
        return {
            "message_id": snapshot.message_id,
            "timestamp": snapshot.timestamp,
            "tracked_file_backups": {
                path: {
                    "backup_file_name": backup.backup_file_name,
                    "version": backup.version,
                    "backup_time": backup.backup_time,
                }
                for path, backup in snapshot.tracked_file_backups.items()
            },
        }

    def _deserialize_snapshot(self, row: Dict[str, Any]) -> Optional[ShallowSnapshot]:
        if not isinstance(row.get("message_id"), str):
            return None
        backups: Dict[str, FileBackup] = {}
        for path, backup_raw in (row.get("tracked_file_backups") or {}).items():
            if not isinstance(backup_raw, dict):
                continue
            backups[str(path)] = FileBackup(
                backup_file_name=backup_raw.get("backup_file_name"),
                version=int(backup_raw.get("version") or 1),
                backup_time=str(backup_raw.get("backup_time") or datetime.now().isoformat()),
            )
        return ShallowSnapshot(
            message_id=row["message_id"],
            tracked_file_backups=backups,
            timestamp=str(row.get("timestamp") or datetime.now().isoformat()),
        )

    # ----------------------------
    # path helpers
    # ----------------------------
    def _normalize_path(self, file_path: Path | str) -> str:
        path = Path(file_path).expanduser()
        if not path.is_absolute():
            path = self.project_path / path
        try:
            rel = path.resolve().relative_to(self.project_path)
            return str(rel).replace("\\", "/")
        except ValueError:
            return str(path).replace("\\", "/")

    def _resolve_file_path(self, tracking_path: str) -> Path:
        if Path(tracking_path).is_absolute():
            return Path(tracking_path).expanduser().resolve()
        return (self.project_path / tracking_path).resolve()

    def _backup_file_name(self, tracking_path: str, version: int) -> str:
        file_name_hash = hashlib.sha256(tracking_path.encode("utf-8")).hexdigest()[:16]
        return f"{file_name_hash}@v{version}"

    def _backup_path(self, backup_file_name: str) -> Path:
        return (self.backup_dir / backup_file_name).resolve()

    # ----------------------------
    # backup operations
    # ----------------------------
    def _create_backup(self, file_path: Path, version: int) -> FileBackup:
        backup_file_name = self._backup_file_name(str(file_path), version)
        backup_path = self._backup_path(backup_file_name)
        backup_path.parent.mkdir(parents=True, exist_ok=True)

        if not file_path.exists():
            return FileBackup(backup_file_name=None, version=version, backup_time=datetime.now().isoformat())

        shutil.copy2(file_path, backup_path)
        return FileBackup(
            backup_file_name=backup_file_name,
            version=version,
            backup_time=datetime.now().isoformat(),
        )

    def _get_latest_backup_version(self, tracking_path: str) -> Optional[FileBackup]:
        latest: Optional[FileBackup] = None
        for snapshot in self._snapshots:
            backup = snapshot.tracked_file_backups.get(tracking_path)
            if backup and (latest is None or backup.version > latest.version):
                latest = backup
        return latest

    def _file_changed_since_backup(self, file_path: Path, backup: FileBackup) -> bool:
        if backup.backup_file_name is None:
            return file_path.exists()
        backup_path = self._backup_path(backup.backup_file_name)
        if not file_path.exists():
            return True
        if not backup_path.exists():
            return True
        try:
            if file_path.stat().st_size != backup_path.stat().st_size:
                return True
            current = file_path.read_text(encoding="utf-8", errors="ignore")
            saved = backup_path.read_text(encoding="utf-8", errors="ignore")
            return current != saved
        except Exception:
            return True

    # ----------------------------
    # public API
    # ----------------------------
    def track_edit(self, file_path: Path | str, message_id: str) -> None:
        """Track a file before it is edited; backup its current content if not already tracked."""
        tracking_path = self._normalize_path(file_path)
        from utils.perf_log import perf_log
        perf_log("[ShallowVersioning] track_edit called", extra={
            "tracking_path": tracking_path,
            "message_id": message_id,
            "already_tracked": tracking_path in self._tracked_files,
            "snapshots_count": len(self._snapshots),
        })
        if tracking_path in self._tracked_files:
            return

        resolved = self._resolve_file_path(tracking_path)
        backup = self._create_backup(resolved, 1)
        self._tracked_files.add(tracking_path)

        # Add to the most recent snapshot if one exists, otherwise create a pre-edit snapshot.
        now = datetime.now().isoformat()
        if self._snapshots:
            latest = self._snapshots[-1]
            latest.tracked_file_backups[tracking_path] = backup
            self._persist_snapshot(latest)
        else:
            snapshot = ShallowSnapshot(
                message_id=message_id,
                tracked_file_backups={tracking_path: backup},
                timestamp=now,
            )
            self._snapshots.append(snapshot)
            self._snapshot_sequence += 1
            self._persist_snapshot(snapshot)

    def make_snapshot(self, message_id: str) -> Dict[str, Any]:
        """Create a snapshot for the given message, backing up changed tracked files."""
        from utils.perf_log import perf_log
        perf_log("[ShallowVersioning] make_snapshot called", extra={
            "message_id": message_id,
            "tracked_files_count": len(self._tracked_files),
            "snapshots_count": len(self._snapshots),
        })
        tracked_file_backups: Dict[str, FileBackup] = {}
        now = datetime.now().isoformat()

        for tracking_path in self._tracked_files:
            resolved = self._resolve_file_path(tracking_path)
            latest_backup = self._get_latest_backup_version(tracking_path)
            next_version = (latest_backup.version + 1) if latest_backup else 1

            if latest_backup and not self._file_changed_since_backup(resolved, latest_backup):
                tracked_file_backups[tracking_path] = latest_backup
                continue

            tracked_file_backups[tracking_path] = self._create_backup(resolved, next_version)

        # Merge in backups from the latest snapshot for any files that may have been added
        # by track_edit during the async window.
        if self._snapshots:
            for tracking_path, backup in self._snapshots[-1].tracked_file_backups.items():
                if tracking_path not in tracked_file_backups:
                    tracked_file_backups[tracking_path] = backup

        snapshot = ShallowSnapshot(
            message_id=message_id,
            tracked_file_backups=tracked_file_backups,
            timestamp=now,
        )
        self._snapshots.append(snapshot)
        self._snapshot_sequence += 1
        self._persist_snapshot(snapshot)

        # Evict old snapshots while keeping the most recent MAX_SNAPSHOTS.
        if len(self._snapshots) > self.MAX_SNAPSHOTS:
            self._snapshots = self._snapshots[-self.MAX_SNAPSHOTS :]

        perf_log("[ShallowVersioning] make_snapshot done", extra={
            "message_id": message_id,
            "tracked_files_count": len(tracked_file_backups),
            "backups": {k: v.backup_file_name for k, v in tracked_file_backups.items()},
        })
        return {
            "message_id": message_id,
            "tracked_files_count": len(tracked_file_backups),
            "timestamp": now,
        }

    def can_restore(self, message_id: str) -> bool:
        return any(snapshot.message_id == message_id for snapshot in self._snapshots)

    def rewind(self, message_id: str) -> Dict[str, Any]:
        """Restore tracked files to the state recorded for the target message."""
        target_snapshot: Optional[ShallowSnapshot] = None
        for snapshot in reversed(self._snapshots):
            if snapshot.message_id == message_id:
                target_snapshot = snapshot
                break
        if not target_snapshot:
            raise ShallowVersioningError(f"未找到消息 {message_id} 对应的快照")

        files_changed: List[str] = []
        for tracking_path in self._tracked_files:
            target_backup = target_snapshot.tracked_file_backups.get(tracking_path)
            if not target_backup:
                # File was not tracked at the target snapshot; leave it unchanged.
                continue

            resolved = self._resolve_file_path(tracking_path)
            if target_backup.backup_file_name is None:
                if resolved.exists():
                    resolved.unlink()
                    files_changed.append(tracking_path)
                continue

            backup_path = self._backup_path(target_backup.backup_file_name)
            if not backup_path.exists():
                continue
            if self._file_changed_since_backup(resolved, target_backup):
                resolved.parent.mkdir(parents=True, exist_ok=True)
                shutil.copy2(backup_path, resolved)
                files_changed.append(tracking_path)

        return {
            "success": True,
            "message_id": message_id,
            "files_changed": files_changed,
        }

    def _read_backup_lines(self, backup: Optional[FileBackup]) -> List[str]:
        """Read lines from a backup entry; empty list for missing/deleted files."""
        if not backup or backup.backup_file_name is None:
            return []
        backup_path = self._backup_path(backup.backup_file_name)
        if not backup_path.exists():
            return []
        try:
            return backup_path.read_text(encoding="utf-8", errors="ignore").splitlines()
        except Exception:
            return []

    def _count_line_changes(self, previous_lines: List[str], target_lines: List[str]) -> tuple[int, int]:
        """Count real insertions/deletions between two line lists using SequenceMatcher."""
        if previous_lines == target_lines:
            return 0, 0
        try:
            import difflib
            sm = difflib.SequenceMatcher(None, previous_lines, target_lines)
        except Exception:
            # Fallback to naive delta if difflib fails.
            if len(previous_lines) > len(target_lines):
                return 0, len(previous_lines) - len(target_lines)
            return len(target_lines) - len(previous_lines), 0

        insertions = 0
        deletions = 0
        for tag, i1, i2, j1, j2 in sm.get_opcodes():
            if tag == "insert":
                insertions += j2 - j1
            elif tag == "delete":
                deletions += i2 - i1
            elif tag == "replace":
                insertions += j2 - j1
                deletions += i2 - i1
        return insertions, deletions

    def _find_snapshot_pair(self, message_id: str):
        """Return (target_snapshot, previous_snapshot) for the given message_id."""
        target_snapshot: Optional[ShallowSnapshot] = None
        target_index = -1
        for idx, snapshot in enumerate(self._snapshots):
            if snapshot.message_id == message_id:
                target_snapshot = snapshot
                target_index = idx
        if not target_snapshot or target_index < 0:
            return None, None
        previous_snapshot = self._snapshots[target_index - 1] if target_index > 0 else None
        return target_snapshot, previous_snapshot

    def get_diff_stats(self, message_id: str) -> Optional[Dict[str, Any]]:
        """Compute insertions/deletions/files_changed between the target snapshot and the previous snapshot."""
        target_snapshot, previous_snapshot = self._find_snapshot_pair(message_id)
        if not target_snapshot:
            return None

        files_changed: List[str] = []
        insertions = 0
        deletions = 0

        for tracking_path in self._tracked_files:
            target_backup = target_snapshot.tracked_file_backups.get(tracking_path)
            if not target_backup:
                continue

            previous_backup = previous_snapshot.tracked_file_backups.get(tracking_path) if previous_snapshot else None
            previous_lines = self._read_backup_lines(previous_backup)
            target_lines = self._read_backup_lines(target_backup)

            if previous_lines == target_lines:
                continue

            files_changed.append(tracking_path)
            ins, dels = self._count_line_changes(previous_lines, target_lines)
            insertions += ins
            deletions += dels

        return {
            "filesChanged": files_changed,
            "insertions": insertions,
            "deletions": deletions,
        }

    def has_any_changes(self, message_id: str) -> bool:
        stats = self.get_diff_stats(message_id)
        if not stats:
            return False
        return bool(stats.get("filesChanged"))

    def get_file_diff_stats(self, message_id: str) -> List[Dict[str, Any]]:
        """Return per-file insertions/deletions/status for a target snapshot vs previous snapshot."""
        target_snapshot, previous_snapshot = self._find_snapshot_pair(message_id)
        if not target_snapshot:
            return []

        result: List[Dict[str, Any]] = []
        for tracking_path in self._tracked_files:
            target_backup = target_snapshot.tracked_file_backups.get(tracking_path)
            if not target_backup:
                continue

            previous_backup = previous_snapshot.tracked_file_backups.get(tracking_path) if previous_snapshot else None
            previous_lines = self._read_backup_lines(previous_backup)
            target_lines = self._read_backup_lines(target_backup)

            if previous_lines == target_lines:
                continue

            insertions, deletions = self._count_line_changes(previous_lines, target_lines)

            resolved = self._resolve_file_path(tracking_path)
            status = "deleted" if target_backup.backup_file_name is None else ("added" if previous_backup is None or previous_backup.backup_file_name is None else "modified")
            result.append({
                "path": tracking_path,
                "status": status,
                "insertions": insertions,
                "deletions": deletions,
            })
        return result

    def get_file_patch_lines(
        self,
        tracking_path: str,
        message_id: str,
        max_lines: int = 600,
    ) -> Dict[str, Any]:
        """Return contextual patch lines between target snapshot and previous snapshot."""
        target_snapshot, previous_snapshot = self._find_snapshot_pair(message_id)
        if not target_snapshot:
            return {"lines": [], "truncated": False}

        target_backup = target_snapshot.tracked_file_backups.get(tracking_path)
        if not target_backup:
            return {"lines": [], "truncated": False}

        previous_backup = previous_snapshot.tracked_file_backups.get(tracking_path) if previous_snapshot else None
        previous_lines = self._read_backup_lines(previous_backup)
        target_lines = self._read_backup_lines(target_backup)

        if previous_lines == target_lines:
            return {"lines": [], "truncated": False}

        try:
            import difflib
            sm = difflib.SequenceMatcher(None, previous_lines, target_lines)
        except Exception:
            return {"lines": [], "truncated": False}

        lines: List[Dict[str, Any]] = []
        truncated = False
        for tag, i1, i2, j1, j2 in sm.get_opcodes():
            if tag == "equal":
                continue
            if tag == "replace" or tag == "delete":
                for line in previous_lines[i1:i2]:
                    if len(lines) >= max_lines:
                        truncated = True
                        break
                    lines.append({"type": "remove", "content": line})
            if tag == "replace" or tag == "insert":
                for line in target_lines[j1:j2]:
                    if len(lines) >= max_lines:
                        truncated = True
                        break
                    lines.append({"type": "add", "content": line})
            if truncated:
                break
        return {"lines": lines, "truncated": truncated}

    def get_snapshot_by_seq(self, seq: int) -> Optional[ShallowSnapshot]:
        """Return the user-message snapshot by 1-based index (skips the initial track_edit snapshot)."""
        user_snapshots = [s for s in self._snapshots if s.message_id != self.conversation_id]
        idx = int(seq) - 1
        if idx < 0 or idx >= len(user_snapshots):
            return None
        return user_snapshots[idx]

    def list_snapshots(self) -> List[Dict[str, Any]]:
        return [self._serialize_snapshot(s) for s in self._snapshots]