zai-cli 0.1.0__py3-none-any.whl

zai/core/undo.py

@@ -0,0 +1,443 @@
+import hashlib
+import json
+import uuid
+from datetime import datetime, timezone
+from pathlib import Path
+
+from .errors import FileError
+from .security import resolve_project_path
+from .storage import atomic_write_json, atomic_write_text, file_lock
+
+MAX_ACTIONS = 50
+MAX_BACKUP_BYTES = 2 * 1024 * 1024
+MAX_TOTAL_BACKUP_BYTES = 10 * 1024 * 1024
+SENSITIVE_NAMES = {
+    ".env",
+    ".env.local",
+    ".env.production",
+    "credentials.json",
+    "secrets.json",
+    "id_rsa",
+    "id_ed25519",
+}
+SENSITIVE_SUFFIXES = {".pem", ".key", ".p12", ".pfx"}
+
+
+def _undo_dir(cwd: str | Path) -> Path:
+    return Path(cwd).resolve() / ".zai" / "undo"
+
+
+def _journal_path(cwd: str | Path) -> Path:
+    return _undo_dir(cwd) / "journal.json"
+
+
+def _redo_path(cwd: str | Path) -> Path:
+    return _undo_dir(cwd) / "redo.json"
+
+
+def _load_actions_unlocked(cwd: str | Path) -> list[dict]:
+    journal = _journal_path(cwd)
+    if not journal.exists():
+        return []
+    try:
+        data = json.loads(journal.read_text(encoding="utf-8"))
+        return data if isinstance(data, list) else []
+    except (OSError, UnicodeError, json.JSONDecodeError):
+        corrupt = journal.with_name(
+            f"{journal.name}.corrupt-"
+            f"{datetime.now().strftime('%Y%m%d-%H%M%S')}"
+        )
+        try:
+            journal.replace(corrupt)
+        except OSError:
+            pass
+        return []
+
+
+def _load_actions(cwd: str | Path) -> list[dict]:
+    with file_lock(_journal_path(cwd)):
+        return _load_actions_unlocked(cwd)
+
+
+def _save_actions_unlocked(cwd: str | Path, actions: list[dict]) -> None:
+    atomic_write_json(_journal_path(cwd), actions, lock=False)
+
+
+def _load_redo_unlocked(cwd: str | Path) -> list[dict]:
+    path = _redo_path(cwd)
+    if not path.exists():
+        return []
+    try:
+        data = json.loads(path.read_text(encoding="utf-8"))
+        return data if isinstance(data, list) else []
+    except (OSError, UnicodeError, json.JSONDecodeError):
+        return []
+
+
+def _save_redo_unlocked(cwd: str | Path, actions: list[dict]) -> None:
+    atomic_write_json(_redo_path(cwd), actions, lock=False)
+
+
+def _relative(cwd: str | Path, path: str | Path) -> str:
+    return str(Path(path).resolve().relative_to(Path(cwd).resolve()))
+
+
+def _hash_file(path: Path) -> str | None:
+    if not path.is_file():
+        return None
+    digest = hashlib.sha256()
+    try:
+        with path.open("rb") as handle:
+            for chunk in iter(lambda: handle.read(1024 * 1024), b""):
+                digest.update(chunk)
+        return digest.hexdigest()
+    except OSError:
+        return None
+
+
+def _is_sensitive(path: Path) -> bool:
+    name = path.name.lower()
+    return (
+        name in SENSITIVE_NAMES
+        or path.suffix.lower() in SENSITIVE_SUFFIXES
+        or name.startswith(".env.")
+        or any(part.lower() in {".ssh", ".aws", ".gnupg"} for part in path.parts)
+    )
+
+
+def _delete_backup(cwd: str | Path, action: dict) -> None:
+    backup = action.get("backup")
+    if backup:
+        (_undo_dir(cwd) / backup).unlink(missing_ok=True)
+    redo_backup = action.get("redo_backup")
+    if redo_backup:
+        (_undo_dir(cwd) / redo_backup).unlink(missing_ok=True)
+
+
+def _action_paths(action: dict) -> list[str]:
+    if action.get("type") == "rename":
+        return [action.get("source", ""), action.get("destination", "")]
+    return [action.get("path", "")]
+
+
+def _summary(action: dict) -> str:
+    action_type = action.get("type")
+    paths = _action_paths(action)
+    if action_type == "restore_file":
+        return f"Restore previous content of {paths[0]}"
+    if action_type == "delete_file":
+        return f"Delete newly created file {paths[0]}"
+    if action_type == "rename":
+        return f"Rename {paths[1]} back to {paths[0]}"
+    if action_type == "remove_folder":
+        return f"Remove newly created folder {paths[0]}"
+    return f"Unknown operation on {', '.join(filter(None, paths))}"
+
+
+def _normalize_action(action: dict) -> dict:
+    normalized = dict(action)
+    normalized.setdefault(
+        "id",
+        hashlib.sha256(
+            json.dumps(action, sort_keys=True).encode("utf-8")
+        ).hexdigest()[:12],
+    )
+    normalized.setdefault("summary", _summary(normalized))
+    normalized.setdefault("paths", list(filter(None, _action_paths(normalized))))
+    return normalized
+
+
+def _prune(cwd: str | Path, actions: list[dict]) -> list[dict]:
+    while len(actions) > MAX_ACTIONS:
+        _delete_backup(cwd, actions.pop(0))
+
+    def backup_size(action: dict) -> int:
+        backup = action.get("backup")
+        if not backup:
+            return 0
+        try:
+            return (_undo_dir(cwd) / backup).stat().st_size
+        except OSError:
+            return 0
+
+    total = sum(backup_size(action) for action in actions)
+    while actions and total > MAX_TOTAL_BACKUP_BYTES:
+        removed = actions.pop(0)
+        total -= backup_size(removed)
+        _delete_backup(cwd, removed)
+    return actions
+
+
+def _append_action(cwd: str | Path, action: dict) -> None:
+    journal = _journal_path(cwd)
+    with file_lock(journal):
+        actions = _load_actions_unlocked(cwd)
+        actions.append(_normalize_action({
+            **action,
+            "id": uuid.uuid4().hex,
+            "created_at": datetime.now(timezone.utc).isoformat(),
+        }))
+        for redo_action in _load_redo_unlocked(cwd):
+            _delete_backup(cwd, redo_action)
+        _save_redo_unlocked(cwd, [])
+        _save_actions_unlocked(cwd, _prune(cwd, actions))
+
+
+def record_file_change(
+    cwd: str,
+    path: str | Path,
+    previous_content: str | None,
+) -> bool:
+    """Record a verified file mutation. Returns False when backup is skipped."""
+    target = Path(path).resolve()
+    action = {
+        "type": "delete_file" if previous_content is None else "restore_file",
+        "path": _relative(cwd, target),
+        "expected_hash": _hash_file(target),
+    }
+    if previous_content is not None:
+        encoded = previous_content.encode("utf-8")
+        if _is_sensitive(target) or len(encoded) > MAX_BACKUP_BYTES:
+            return False
+        backup_name = f"{uuid.uuid4().hex}.txt"
+        backup_path = _undo_dir(cwd) / backup_name
+        atomic_write_text(backup_path, previous_content, lock=False)
+        action["backup"] = backup_name
+        action["backup_hash"] = hashlib.sha256(encoded).hexdigest()
+    _append_action(cwd, action)
+    return True
+
+
+def record_rename(cwd: str, source: str | Path, destination: str | Path) -> None:
+    destination_path = Path(destination).resolve()
+    _append_action(cwd, {
+        "type": "rename",
+        "source": _relative(cwd, source),
+        "destination": _relative(cwd, destination_path),
+        "expected_hash": _hash_file(destination_path),
+    })
+
+
+def record_created_folder(cwd: str, path: str | Path) -> None:
+    _append_action(cwd, {
+        "type": "remove_folder",
+        "path": _relative(cwd, path),
+    })
+
+
+def _conflict(path: Path, expected_hash: str | None) -> bool:
+    if expected_hash is None:
+        return False
+    return _hash_file(path) != expected_hash
+
+
+def list_actions(cwd: str) -> list[dict]:
+    """Return inspectable undo actions, newest first."""
+    return [_normalize_action(item) for item in reversed(_load_actions(cwd))]
+
+
+def get_action(cwd: str, identifier: str) -> dict | None:
+    matches = [
+        item
+        for item in list_actions(cwd)
+        if item["id"] == identifier or item["id"].startswith(identifier)
+    ]
+    return matches[0] if len(matches) == 1 else None
+
+
+def _write_redo_backup(cwd: str, content: str) -> tuple[str, str]:
+    encoded = content.encode("utf-8")
+    backup_name = f"redo-{uuid.uuid4().hex}.txt"
+    atomic_write_text(_undo_dir(cwd) / backup_name, content, lock=False)
+    return backup_name, hashlib.sha256(encoded).hexdigest()
+
+
+def _undo_one(cwd: str, action: dict) -> tuple[str, dict | None]:
+    action_type = action.get("type")
+    redo_action = dict(action)
+    try:
+        if action_type == "restore_file":
+            path = resolve_project_path(cwd, action["path"])
+            backup = _undo_dir(cwd) / action["backup"]
+            if not backup.is_file():
+                return "Cannot undo: backup file is missing.", None
+            if _conflict(path, action.get("expected_hash")):
+                return (
+                    f"Cannot undo: {path.name} changed after the AI edit. "
+                    "Current content was preserved."
+                ), None
+            current = path.read_text(encoding="utf-8")
+            redo_backup, redo_hash = _write_redo_backup(cwd, current)
+            previous = backup.read_text(encoding="utf-8")
+            if hashlib.sha256(previous.encode("utf-8")).hexdigest() != action.get(
+                "backup_hash"
+            ):
+                _delete_backup(cwd, {"redo_backup": redo_backup})
+                return "Cannot undo: backup integrity check failed.", None
+            atomic_write_text(path, previous, mode=0o644, lock=False)
+            redo_action.update({
+                "redo_backup": redo_backup,
+                "redo_hash": redo_hash,
+                "undo_hash": _hash_file(path),
+            })
+            message = f"Restored: {path.name}"
+        elif action_type == "delete_file":
+            path = resolve_project_path(cwd, action["path"])
+            if path.exists() and not path.is_file():
+                return f"Cannot undo: not a file: {path.name}", None
+            if path.exists() and _conflict(path, action.get("expected_hash")):
+                return (
+                    f"Cannot undo: {path.name} changed after creation. "
+                    "Current content was preserved."
+                ), None
+            if path.exists():
+                current = path.read_text(encoding="utf-8")
+                redo_backup, redo_hash = _write_redo_backup(cwd, current)
+                redo_action.update({
+                    "redo_backup": redo_backup,
+                    "redo_hash": redo_hash,
+                })
+                path.unlink()
+            message = f"Deleted new file: {path.name}"
+        elif action_type == "rename":
+            source = resolve_project_path(cwd, action["source"])
+            destination = resolve_project_path(cwd, action["destination"])
+            if not destination.exists():
+                return f"Cannot undo rename: path missing: {destination.name}", None
+            if source.exists():
+                return f"Cannot undo rename: destination exists: {source.name}", None
+            if destination.is_file() and _conflict(
+                destination,
+                action.get("expected_hash"),
+            ):
+                return (
+                    f"Cannot undo rename: {destination.name} changed after rename."
+                ), None
+            source.parent.mkdir(parents=True, exist_ok=True)
+            destination.rename(source)
+            redo_action["undo_hash"] = _hash_file(source)
+            message = f"Restored name: {destination.name} -> {source.name}"
+        elif action_type == "remove_folder":
+            path = resolve_project_path(cwd, action["path"])
+            if path.exists() and not path.is_dir():
+                return f"Cannot undo: not a folder: {path.name}", None
+            if path.exists():
+                try:
+                    path.rmdir()
+                except OSError:
+                    return f"Cannot undo: folder is not empty: {path.name}", None
+            message = f"Deleted new folder: {path.name}"
+        else:
+            return "Cannot undo: unknown journal action.", None
+    except (FileError, OSError, UnicodeError) as error:
+        return f"Cannot undo: {error}", None
+    return message, redo_action
+
+
+def undo_action(cwd: str, identifier: str | None = None) -> str:
+    """Undo the latest action, or all actions through a selected ID."""
+    journal = _journal_path(cwd)
+    with file_lock(journal):
+        actions = _load_actions_unlocked(cwd)
+        if not actions:
+            return "Nothing to undo."
+        normalized = [_normalize_action(item) for item in actions]
+        target_index = len(actions) - 1
+        if identifier:
+            matches = [
+                index
+                for index, item in enumerate(normalized)
+                if item["id"] == identifier or item["id"].startswith(identifier)
+            ]
+            if len(matches) != 1:
+                return "Cannot undo: action ID not found or ambiguous."
+            target_index = matches[0]
+        redo_actions = _load_redo_unlocked(cwd)
+        messages: list[str] = []
+        while len(actions) - 1 >= target_index:
+            action = normalized[len(actions) - 1]
+            message, redo_action = _undo_one(cwd, action)
+            if redo_action is None:
+                _save_actions_unlocked(cwd, actions)
+                _save_redo_unlocked(cwd, redo_actions)
+                return "\n".join(messages + [message])
+            actions.pop()
+            redo_actions.append(redo_action)
+            messages.append(message)
+        _save_actions_unlocked(cwd, actions)
+        _save_redo_unlocked(cwd, redo_actions)
+        return "\n".join(messages)
+
+
+def undo_last_action(cwd: str) -> str:
+    return undo_action(cwd)
+
+
+def redo_last_action(cwd: str) -> str:
+    journal = _journal_path(cwd)
+    with file_lock(journal):
+        redo_actions = _load_redo_unlocked(cwd)
+        if not redo_actions:
+            return "Nothing to redo."
+        action = redo_actions[-1]
+        action_type = action.get("type")
+        try:
+            if action_type in {"restore_file", "delete_file"}:
+                path = resolve_project_path(cwd, action["path"])
+                expected = action.get("undo_hash")
+                if path.exists() and _conflict(path, expected):
+                    return f"Cannot redo: {path.name} changed after undo."
+                backup = _undo_dir(cwd) / action.get("redo_backup", "")
+                if not backup.is_file():
+                    return "Cannot redo: redo backup is missing."
+                content = backup.read_text(encoding="utf-8")
+                if hashlib.sha256(content.encode("utf-8")).hexdigest() != action.get(
+                    "redo_hash"
+                ):
+                    return "Cannot redo: redo backup integrity check failed."
+                atomic_write_text(path, content, mode=0o644, lock=False)
+                action["expected_hash"] = _hash_file(path)
+                message = f"Redone file change: {path.name}"
+            elif action_type == "rename":
+                source = resolve_project_path(cwd, action["source"])
+                destination = resolve_project_path(cwd, action["destination"])
+                if not source.exists() or destination.exists():
+                    return "Cannot redo rename: source missing or destination exists."
+                if source.is_file() and _conflict(source, action.get("undo_hash")):
+                    return f"Cannot redo rename: {source.name} changed after undo."
+                destination.parent.mkdir(parents=True, exist_ok=True)
+                source.rename(destination)
+                action["expected_hash"] = _hash_file(destination)
+                message = f"Redone rename: {source.name} -> {destination.name}"
+            elif action_type == "remove_folder":
+                path = resolve_project_path(cwd, action["path"])
+                if path.exists():
+                    return f"Cannot redo: path already exists: {path.name}"
+                path.mkdir(parents=True)
+                message = f"Redone folder creation: {path.name}"
+            else:
+                return "Cannot redo: unknown journal action."
+        except (FileError, OSError, UnicodeError) as error:
+            return f"Cannot redo: {error}"
+        redo_actions.pop()
+        _delete_backup(cwd, {"redo_backup": action.get("redo_backup")})
+        action.pop("redo_backup", None)
+        action.pop("redo_hash", None)
+        action.pop("undo_hash", None)
+        actions = _load_actions_unlocked(cwd)
+        actions.append(action)
+        _save_actions_unlocked(cwd, _prune(cwd, actions))
+        _save_redo_unlocked(cwd, redo_actions)
+        return message
+
+
+def clear_actions(cwd: str) -> int:
+    journal = _journal_path(cwd)
+    with file_lock(journal):
+        actions = _load_actions_unlocked(cwd)
+        redo_actions = _load_redo_unlocked(cwd)
+        for action in actions + redo_actions:
+            _delete_backup(cwd, action)
+        _save_actions_unlocked(cwd, [])
+        _save_redo_unlocked(cwd, [])
+        return len(actions) + len(redo_actions)

zai/core/watch.py

@@ -0,0 +1,80 @@
+import os
+import time
+from pathlib import Path
+from rich.console import Console
+
+console = Console()
+
+_observer = None
+_watching = False
+
+SKIP_NAMES = {'.git', '__pycache__', 'node_modules', '.venv', 'dist', 'build', '.mypy_cache'}
+SKIP_EXT = {'.pyc', '.pyo', '.pyd', '.swp', '.tmp', '.log'}
+
+
+def _should_skip(path: str) -> bool:
+    p = Path(path)
+    if p.suffix in SKIP_EXT:
+        return True
+    for part in p.parts:
+        if part in SKIP_NAMES:
+            return True
+    return False
+
+
+def start_watch(cwd: str, callback) -> bool:
+    """
+    Start watching `cwd` for file changes.
+    `callback(filename)` is called when a relevant file is modified.
+    Returns True if started, False if watchdog not installed.
+    """
+    global _observer, _watching
+
+    try:
+        from watchdog.observers import Observer
+        from watchdog.events import FileSystemEventHandler
+    except ImportError:
+        console.print("[yellow]Install watchdog to use watch mode: pip install watchdog[/yellow]")
+        return False
+
+    class Handler(FileSystemEventHandler):
+        def __init__(self):
+            self._cooldown: dict = {}
+
+        def on_modified(self, event):
+            if event.is_directory:
+                return
+            src = event.src_path
+            if _should_skip(src):
+                return
+            now = time.time()
+            if now - self._cooldown.get(src, 0) < 3:
+                return
+            self._cooldown[src] = now
+            fname = os.path.basename(src)
+            console.print(f"\n[dim]~ {fname} changed[/dim]")
+            callback(fname)
+
+    if _observer:
+        _observer.stop()
+        _observer.join()
+
+    _observer = Observer()
+    _observer.schedule(Handler(), cwd, recursive=True)
+    _observer.start()
+    _watching = True
+    return True
+
+
+def stop_watch():
+    """Stop the file watcher."""
+    global _observer, _watching
+    if _observer:
+        _observer.stop()
+        _observer.join()
+        _observer = None
+    _watching = False
+
+
+def is_watching() -> bool:
+    return _watching