aria-code 0.52.0__py3-none-any.whl

aria_code/runtime/checkpoints.py

@@ -0,0 +1,432 @@
+"""Durable, conflict-aware file checkpoints for agent write operations."""
+
+from __future__ import annotations
+
+import hashlib
+import json
+import os
+import sqlite3
+import stat
+import time
+import uuid
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Any, Dict, Iterable, Optional
+
+from .run_store import RunStore
+
+
+def _sha256(content: str) -> str:
+    return hashlib.sha256(content.encode("utf-8", errors="replace")).hexdigest()
+
+
+def _json_dumps(value: Any) -> str:
+    return json.dumps(value, ensure_ascii=False, default=str, separators=(",", ":"))
+
+
+def _json_loads(value: str | None) -> Dict[str, Any]:
+    if not value:
+        return {}
+    try:
+        loaded = json.loads(value)
+        return loaded if isinstance(loaded, dict) else {"value": loaded}
+    except (TypeError, json.JSONDecodeError):
+        return {}
+
+
+@dataclass(frozen=True)
+class CheckpointFile:
+    checkpoint_id: str
+    path: str
+    existed_before: bool
+    before_content: str
+    after_content: str
+    before_hash: str
+    after_hash: str
+    before_mode: Optional[int]
+
+
+@dataclass(frozen=True)
+class CheckpointRecord:
+    checkpoint_id: str
+    sequence: int
+    run_id: Optional[str]
+    session_id: str
+    source: str
+    label: str
+    status: str
+    created_at: float
+    restored_at: Optional[float]
+    metadata: Dict[str, Any] = field(default_factory=dict)
+    files: tuple[CheckpointFile, ...] = ()
+
+
+@dataclass(frozen=True)
+class RestoreResult:
+    checkpoint_ids: tuple[str, ...]
+    run_id: Optional[str]
+    restored_paths: tuple[str, ...]
+
+
+class CheckpointNotFoundError(KeyError):
+    """Raised when a checkpoint or checkpoint group cannot be found."""
+
+
+class CheckpointConflictError(RuntimeError):
+    """Raised when a file changed after the checkpoint was recorded."""
+
+
+class CheckpointStore:
+    """Store file pre-images and restore them only when hashes still match."""
+
+    def __init__(self, database_path: Path | str | None = None) -> None:
+        self.run_store = RunStore(database_path)
+        self.database_path = self.run_store.database_path
+        self._initialize()
+
+    def _connect(self) -> sqlite3.Connection:
+        connection = sqlite3.connect(str(self.database_path), timeout=10.0)
+        connection.row_factory = sqlite3.Row
+        connection.execute("PRAGMA foreign_keys = ON")
+        connection.execute("PRAGMA busy_timeout = 10000")
+        return connection
+
+    def _initialize(self) -> None:
+        with self._connect() as connection:
+            connection.executescript(
+                """
+                CREATE TABLE IF NOT EXISTS checkpoints (
+                    sequence INTEGER PRIMARY KEY AUTOINCREMENT,
+                    checkpoint_id TEXT NOT NULL UNIQUE,
+                    run_id TEXT,
+                    session_id TEXT NOT NULL DEFAULT '',
+                    source TEXT NOT NULL,
+                    label TEXT NOT NULL DEFAULT '',
+                    status TEXT NOT NULL DEFAULT 'active',
+                    created_at REAL NOT NULL,
+                    restored_at REAL,
+                    metadata_json TEXT NOT NULL DEFAULT '{}',
+                    FOREIGN KEY(run_id) REFERENCES runs(run_id)
+                );
+
+                CREATE TABLE IF NOT EXISTS checkpoint_files (
+                    checkpoint_id TEXT NOT NULL,
+                    path TEXT NOT NULL,
+                    existed_before INTEGER NOT NULL,
+                    before_content TEXT NOT NULL,
+                    after_content TEXT NOT NULL,
+                    before_hash TEXT NOT NULL,
+                    after_hash TEXT NOT NULL,
+                    before_mode INTEGER,
+                    PRIMARY KEY(checkpoint_id, path),
+                    FOREIGN KEY(checkpoint_id)
+                        REFERENCES checkpoints(checkpoint_id) ON DELETE CASCADE
+                );
+
+                CREATE INDEX IF NOT EXISTS idx_checkpoints_run_sequence
+                    ON checkpoints(run_id, sequence DESC);
+                CREATE INDEX IF NOT EXISTS idx_checkpoints_session_sequence
+                    ON checkpoints(session_id, sequence DESC);
+                CREATE INDEX IF NOT EXISTS idx_checkpoints_status_sequence
+                    ON checkpoints(status, sequence DESC);
+                """
+            )
+
+    @staticmethod
+    def _row_to_file(row: sqlite3.Row) -> CheckpointFile:
+        return CheckpointFile(
+            checkpoint_id=row["checkpoint_id"],
+            path=row["path"],
+            existed_before=bool(row["existed_before"]),
+            before_content=row["before_content"],
+            after_content=row["after_content"],
+            before_hash=row["before_hash"],
+            after_hash=row["after_hash"],
+            before_mode=int(row["before_mode"]) if row["before_mode"] is not None else None,
+        )
+
+    def _row_to_checkpoint(
+        self,
+        connection: sqlite3.Connection,
+        row: sqlite3.Row,
+    ) -> CheckpointRecord:
+        file_rows = connection.execute(
+            "SELECT * FROM checkpoint_files WHERE checkpoint_id = ? ORDER BY path",
+            (row["checkpoint_id"],),
+        ).fetchall()
+        return CheckpointRecord(
+            checkpoint_id=row["checkpoint_id"],
+            sequence=int(row["sequence"]),
+            run_id=row["run_id"],
+            session_id=row["session_id"],
+            source=row["source"],
+            label=row["label"],
+            status=row["status"],
+            created_at=float(row["created_at"]),
+            restored_at=float(row["restored_at"]) if row["restored_at"] is not None else None,
+            metadata=_json_loads(row["metadata_json"]),
+            files=tuple(self._row_to_file(file_row) for file_row in file_rows),
+        )
+
+    def record_change(
+        self,
+        *,
+        path: Path | str,
+        before_content: str,
+        after_content: str,
+        existed_before: bool,
+        source: str,
+        run_id: str | None = None,
+        session_id: str = "",
+        before_mode: int | None = None,
+        label: str = "",
+        metadata: Dict[str, Any] | None = None,
+    ) -> CheckpointRecord:
+        target = Path(path).expanduser().resolve()
+        if not target.is_file():
+            raise CheckpointConflictError(f"Written file is missing: {target}")
+        current = target.read_text(encoding="utf-8", errors="replace")
+        after_hash = _sha256(after_content)
+        if _sha256(current) != after_hash:
+            raise CheckpointConflictError(
+                f"File changed before checkpoint could be recorded: {target}"
+            )
+        if run_id and self.run_store.get_run(run_id) is None:
+            raise CheckpointNotFoundError(f"Unknown run id: {run_id}")
+
+        checkpoint_id = uuid.uuid4().hex[:16]
+        created_at = time.time()
+        mode = before_mode
+        if existed_before and mode is None:
+            mode = stat.S_IMODE(target.stat().st_mode)
+        with self._connect() as connection:
+            connection.execute("BEGIN IMMEDIATE")
+            connection.execute(
+                """INSERT INTO checkpoints
+                   (checkpoint_id, run_id, session_id, source, label, status,
+                    created_at, restored_at, metadata_json)
+                   VALUES (?, ?, ?, ?, ?, 'active', ?, NULL, ?)""",
+                (
+                    checkpoint_id,
+                    run_id,
+                    session_id,
+                    source,
+                    label,
+                    created_at,
+                    _json_dumps(metadata or {}),
+                ),
+            )
+            connection.execute(
+                """INSERT INTO checkpoint_files
+                   (checkpoint_id, path, existed_before, before_content,
+                    after_content, before_hash, after_hash, before_mode)
+                   VALUES (?, ?, ?, ?, ?, ?, ?, ?)""",
+                (
+                    checkpoint_id,
+                    str(target),
+                    int(existed_before),
+                    before_content,
+                    after_content,
+                    _sha256(before_content),
+                    after_hash,
+                    mode,
+                ),
+            )
+            row = connection.execute(
+                "SELECT * FROM checkpoints WHERE checkpoint_id = ?",
+                (checkpoint_id,),
+            ).fetchone()
+
+        if run_id:
+            self.run_store.append_event(run_id, "checkpoint_created", {
+                "checkpoint_id": checkpoint_id,
+                "path": str(target),
+                "source": source,
+            })
+        assert row is not None
+        with self._connect() as connection:
+            return self._row_to_checkpoint(connection, row)
+
+    def get(self, checkpoint_id: str) -> Optional[CheckpointRecord]:
+        key = (checkpoint_id or "").strip()
+        with self._connect() as connection:
+            row = connection.execute(
+                "SELECT * FROM checkpoints WHERE checkpoint_id = ?",
+                (key,),
+            ).fetchone()
+            if row is None and key:
+                rows = connection.execute(
+                    "SELECT * FROM checkpoints WHERE checkpoint_id LIKE ? ORDER BY sequence DESC LIMIT 2",
+                    (f"{key}%",),
+                ).fetchall()
+                row = rows[0] if len(rows) == 1 else None
+            return self._row_to_checkpoint(connection, row) if row is not None else None
+
+    def list(
+        self,
+        *,
+        run_id: str | None = None,
+        session_id: str | None = None,
+        status: str | None = "active",
+        limit: int = 50,
+    ) -> list[CheckpointRecord]:
+        clauses = []
+        values: list[Any] = []
+        if run_id:
+            clauses.append("run_id = ?")
+            values.append(run_id)
+        if session_id:
+            clauses.append("session_id = ?")
+            values.append(session_id)
+        if status:
+            clauses.append("status = ?")
+            values.append(status)
+        where = f" WHERE {' AND '.join(clauses)}" if clauses else ""
+        values.append(max(1, int(limit)))
+        with self._connect() as connection:
+            rows = connection.execute(
+                f"SELECT * FROM checkpoints{where} ORDER BY sequence DESC LIMIT ?",
+                values,
+            ).fetchall()
+            return [self._row_to_checkpoint(connection, row) for row in rows]
+
+    def latest_run_id(self, *, session_id: str | None = None) -> Optional[str]:
+        clauses = ["status = 'active'", "run_id IS NOT NULL"]
+        values: list[Any] = []
+        if session_id:
+            clauses.append("session_id = ?")
+            values.append(session_id)
+        with self._connect() as connection:
+            row = connection.execute(
+                f"SELECT run_id FROM checkpoints WHERE {' AND '.join(clauses)} "
+                "ORDER BY sequence DESC LIMIT 1",
+                values,
+            ).fetchone()
+        return row["run_id"] if row is not None else None
+
+    @staticmethod
+    def _snapshot(path: Path) -> tuple[bool, str, Optional[int]]:
+        if not path.exists():
+            return False, "", None
+        if not path.is_file():
+            raise CheckpointConflictError(f"Checkpoint path is not a file: {path}")
+        return (
+            True,
+            path.read_text(encoding="utf-8", errors="replace"),
+            stat.S_IMODE(path.stat().st_mode),
+        )
+
+    @staticmethod
+    def _write_atomic(path: Path, content: str, mode: int | None) -> None:
+        path.parent.mkdir(parents=True, exist_ok=True)
+        temporary = path.with_name(f".{path.name}.aria-rewind-{uuid.uuid4().hex[:8]}")
+        try:
+            temporary.write_text(content, encoding="utf-8")
+            if mode is not None:
+                temporary.chmod(mode)
+            os.replace(temporary, path)
+        finally:
+            if temporary.exists():
+                temporary.unlink()
+
+    def _restore(self, checkpoints: Iterable[CheckpointRecord]) -> RestoreResult:
+        records = list(checkpoints)
+        if not records:
+            raise CheckpointNotFoundError("No active checkpoints to restore")
+        files = [file for record in records for file in record.files]
+        if not files:
+            raise CheckpointNotFoundError("Checkpoint contains no files")
+
+        simulated: Dict[str, tuple[bool, str]] = {}
+        backups: Dict[str, tuple[bool, str, Optional[int]]] = {}
+        for file in files:
+            path = Path(file.path)
+            if file.path not in simulated:
+                exists, content, mode = self._snapshot(path)
+                backups[file.path] = (exists, content, mode)
+                simulated[file.path] = (exists, _sha256(content))
+            exists, current_hash = simulated[file.path]
+            if not exists or current_hash != file.after_hash:
+                raise CheckpointConflictError(
+                    f"Refusing to overwrite a file changed after the checkpoint: {path}"
+                )
+            simulated[file.path] = (file.existed_before, file.before_hash)
+
+        restored_paths: list[str] = []
+        try:
+            for file in files:
+                path = Path(file.path)
+                exists, current, _ = self._snapshot(path)
+                if not exists or _sha256(current) != file.after_hash:
+                    raise CheckpointConflictError(
+                        f"File changed while rewind was running: {path}"
+                    )
+                if file.existed_before:
+                    self._write_atomic(path, file.before_content, file.before_mode)
+                else:
+                    path.unlink()
+                if file.path not in restored_paths:
+                    restored_paths.append(file.path)
+        except Exception:
+            for path_text, (existed, content, mode) in backups.items():
+                path = Path(path_text)
+                if existed:
+                    self._write_atomic(path, content, mode)
+                elif path.exists() and path.is_file():
+                    path.unlink()
+            raise
+
+        restored_at = time.time()
+        checkpoint_ids = tuple(record.checkpoint_id for record in records)
+        with self._connect() as connection:
+            connection.execute("BEGIN IMMEDIATE")
+            connection.executemany(
+                "UPDATE checkpoints SET status = 'restored', restored_at = ? WHERE checkpoint_id = ?",
+                [(restored_at, checkpoint_id) for checkpoint_id in checkpoint_ids],
+            )
+
+        run_id = records[0].run_id if all(
+            record.run_id == records[0].run_id for record in records
+        ) else None
+        if run_id:
+            self.run_store.append_event(run_id, "checkpoint_restored", {
+                "checkpoint_ids": list(checkpoint_ids),
+                "paths": restored_paths,
+            })
+        return RestoreResult(checkpoint_ids, run_id, tuple(restored_paths))
+
+    def restore_checkpoint(self, checkpoint_id: str) -> RestoreResult:
+        record = self.get(checkpoint_id)
+        if record is None or record.status != "active":
+            raise CheckpointNotFoundError(checkpoint_id)
+        return self._restore([record])
+
+    def restore_run(self, run_id: str) -> RestoreResult:
+        records = self.list(run_id=run_id, status="active", limit=10_000)
+        if not records:
+            raise CheckpointNotFoundError(run_id)
+        return self._restore(records)
+
+    def restore_latest_run(self, *, session_id: str | None = None) -> RestoreResult:
+        run_id = self.latest_run_id(session_id=session_id)
+        if not run_id:
+            raise CheckpointNotFoundError("No active run checkpoint found")
+        return self.restore_run(run_id)
+
+    def restore_latest(self, *, session_id: str | None = None) -> RestoreResult:
+        clauses = ["status = 'active'"]
+        values: list[Any] = []
+        if session_id:
+            clauses.append("session_id = ?")
+            values.append(session_id)
+        with self._connect() as connection:
+            row = connection.execute(
+                f"SELECT checkpoint_id, run_id FROM checkpoints "
+                f"WHERE {' AND '.join(clauses)} ORDER BY sequence DESC LIMIT 1",
+                values,
+            ).fetchone()
+        if row is None:
+            raise CheckpointNotFoundError("No active checkpoint found")
+        if row["run_id"]:
+            return self.restore_run(row["run_id"])
+        return self.restore_checkpoint(row["checkpoint_id"])

aria_code/runtime/crash_report.py

@@ -0,0 +1,174 @@
+"""崩溃记录：本地优先，不向任何第三方发送。
+
+问题：CLI 未捕获异常时只会打一个裸 traceback 就退出，什么都不留。用户来报
+bug 时能提供的往往只有"它崩了"——版本、provider、走到哪一步、当时的配置
+全部丢失，只能靠来回追问重建现场。
+
+为什么不直接接 Sentry：aria-code 是 local-first 的终端工具，用户的堆栈里可能
+带着文件路径、项目名、甚至提示词片段。默认把这些发到第三方服务器，跟这个
+产品的定位相悖。这里的做法是**先在本地留全**，用户愿意时再自己把文件贴出来；
+需要 Sentry 的部署可以自行在外层包一层，本模块不阻止。
+
+写入 ~/.aria-code/crashes/（随 aria_home() 走，老用户仍落在 ~/.arthera/）。
+
+脱敏是**尽力而为**，不是保证：环境变量里的密钥、明显的 key 形态会被打码，
+但堆栈的局部变量里可能藏着任何东西。所以文件默认只在本地，且提示用户贴出前
+自己看一眼。把这条说清楚，比假装"已完全脱敏"要诚实。
+"""
+
+from __future__ import annotations
+
+import json
+import os
+import platform
+import re
+import sys
+import traceback
+from datetime import datetime, timezone
+from pathlib import Path
+from typing import Any, Dict, List, Optional
+
+__all__ = ["write_crash_report", "install_excepthook", "recent_crashes", "redact"]
+
+_MAX_KEPT = 20
+
+# 环境变量名里出现这些片段，值一律打码
+_SECRET_HINTS = ("KEY", "TOKEN", "SECRET", "PASSWORD", "CREDENTIAL", "AUTH")
+
+# 值本身长得像密钥的，即使变量名无辜也打码
+_SECRET_VALUE_RE = re.compile(
+    r"\b(?:sk-[A-Za-z0-9_-]{16,}|gh[pousr]_[A-Za-z0-9]{20,}|xox[baprs]-[A-Za-z0-9-]{10,}"
+    r"|AIza[A-Za-z0-9_-]{30,}|[A-Za-z0-9_-]{32,}\.[A-Za-z0-9_-]{16,}\.[A-Za-z0-9_-]{16,})\b"
+)
+
+
+def redact(text: str) -> str:
+    """给明显的密钥形态打码。尽力而为——见模块文档。"""
+    return _SECRET_VALUE_RE.sub("<redacted>", text or "")
+
+
+def _crash_dir() -> Path:
+    try:
+        from packages.aria_core.paths import aria_home
+        root = aria_home()
+    except Exception:
+        root = Path.home() / ".aria-code"
+    return root / "crashes"
+
+
+def _safe_env() -> Dict[str, str]:
+    """只收跟 aria 相关的环境变量，且对疑似密钥打码。
+
+    不收全量环境：那里面有用户的 PATH、代理地址、其它项目的凭证，跟这次崩溃
+    多半无关，收进来只是扩大暴露面。
+    """
+    out: Dict[str, str] = {}
+    for key, value in os.environ.items():
+        if not any(p in key.upper() for p in ("ARIA", "ARTHERA", "OLLAMA", "LMSTUDIO")):
+            continue
+        if any(h in key.upper() for h in _SECRET_HINTS):
+            out[key] = f"<set:{len(value)}chars>" if value else "<empty>"
+        else:
+            out[key] = redact(value)[:200]
+    return out
+
+
+def write_crash_report(
+    exc: BaseException,
+    *,
+    context: Optional[Dict[str, Any]] = None,
+) -> Optional[Path]:
+    """把一次崩溃写成本地 JSON。返回文件路径；写入失败返回 None。
+
+    写入失败绝不抛异常——崩溃记录器自己再崩一次，只会把原始错误从屏幕上
+    冲掉，让排查更难。
+    """
+    try:
+        crash_dir = _crash_dir()
+        crash_dir.mkdir(parents=True, exist_ok=True)
+
+        stamp = datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%SZ")
+        path = crash_dir / f"crash-{stamp}.json"
+
+        try:
+            from aria_code._version import __version__ as _ver
+        except Exception:
+            _ver = "unknown"
+
+        payload = {
+            "timestamp": datetime.now(timezone.utc).isoformat(),
+            "version": _ver,
+            "python": sys.version.split()[0],
+            "platform": platform.platform(),
+            "error_type": type(exc).__name__,
+            "error": redact(str(exc))[:1000],
+            "traceback": redact("".join(traceback.format_exception(exc)))[-8000:],
+            "argv": [redact(a) for a in sys.argv[:12]],
+            "env": _safe_env(),
+            "context": {k: redact(str(v))[:500] for k, v in (context or {}).items()},
+        }
+        path.write_text(json.dumps(payload, ensure_ascii=False, indent=2), encoding="utf-8")
+        _prune(crash_dir)
+        return path
+    except Exception:
+        return None
+
+
+def _prune(crash_dir: Path, keep: int = _MAX_KEPT) -> None:
+    """只保留最近 N 份。崩溃记录是排查素材，不是需要永久保存的资产；
+    无上限增长的日志目录本身就会变成一个问题。"""
+    try:
+        files = sorted(crash_dir.glob("crash-*.json"), key=lambda p: p.stat().st_mtime)
+        for stale in files[:-keep]:
+            stale.unlink(missing_ok=True)
+    except Exception:
+        pass
+
+
+def recent_crashes(limit: int = 5) -> List[Dict[str, Any]]:
+    """最近几次崩溃的摘要，供 /doctor 之类的命令展示。"""
+    out: List[Dict[str, Any]] = []
+    try:
+        files = sorted(_crash_dir().glob("crash-*.json"),
+                       key=lambda p: p.stat().st_mtime, reverse=True)
+        for path in files[:limit]:
+            try:
+                data = json.loads(path.read_text(encoding="utf-8"))
+                out.append({
+                    "file": str(path),
+                    "timestamp": data.get("timestamp", ""),
+                    "error_type": data.get("error_type", ""),
+                    "error": (data.get("error", "") or "")[:160],
+                    "version": data.get("version", ""),
+                })
+            except Exception:
+                continue
+    except Exception:
+        pass
+    return out
+
+
+def install_excepthook(*, context_provider=None) -> None:
+    """接管未捕获异常：先落盘，再按原样打印。
+
+    **保留原有输出**——用户和现有的排查习惯都依赖屏幕上那个 traceback，
+    悄悄换掉只会让人以为程序静默失败了。这里只是在它之前多留一份记录。
+
+    KeyboardInterrupt 不记：那是用户主动中断，不是故障。
+    """
+    previous = sys.excepthook
+
+    def _hook(exc_type, exc_value, exc_tb):
+        if not issubclass(exc_type, KeyboardInterrupt):
+            ctx = {}
+            if context_provider is not None:
+                try:
+                    ctx = context_provider() or {}
+                except Exception:
+                    ctx = {}
+            path = write_crash_report(exc_value, context=ctx)
+            if path is not None:
+                print(f"\n[crash report] {path}", file=sys.stderr)
+        previous(exc_type, exc_value, exc_tb)
+
+    sys.excepthook = _hook