feat: add workflow stage snapshots and rollback entrypoint

Introduce a dedicated stage snapshot manager (src/app/snapshots.py) to capture and restore workspace checkpoints with indexed metadata.

Add CLI controls for snapshot operations: --enable-snapshots, --snapshot-dir, and --rollback-to-snapshot. Workflow now supports snapshot restoration as an explicit early-exit recovery path and captures snapshots after initialization, post-modification, and post-repack stages.

Add regression tests for CLI parsing, snapshot capture/restore/list behavior, rollback flow, and snapshot capture invocation order. Update CN/EN README argument tables accordingly.

Verification:

- .venv/bin/python -m ruff check src/app/cli.py src/app/workflow.py src/app/snapshots.py tests/test_cli.py tests/test_workflow.py tests/test_snapshots.py README.md README_EN.md

- .venv/bin/python -m mypy --config-file mypy-curated.ini

- .venv/bin/python -m pytest -q tests/test_workflow.py tests/test_snapshots.py

- .venv/bin/python -m pytest -q

Co-authored-by: OpenAI Codex <noreply@openai.com>

Author: toraidl

README.md

@@ -155,6 +155,9 @@ sudo python3 main.py --stock <底包路径> --port <移植包路径> --pack-type
 | `--skip-preflight` | 跳过预检阶段（不建议） | `false` |
 | `--preflight-strict` | 将风险项也视为失败项（用于严格阻断） | `false` |
 | `--preflight-report` | 预检 JSON 报告输出路径 | `build/preflight-report.json` |
+| `--enable-snapshots` | 在关键阶段保存工作目录快照 | `false` |
+| `--snapshot-dir` | 快照目录（未设置时使用 `<work-dir>/snapshots`） | `null` |
+| `--rollback-to-snapshot` | 从指定快照恢复目标工作目录并退出 | `null` |
 
 ---
 

README_EN.md

@@ -157,6 +157,9 @@ sudo python3 main.py --stock <path_to_stock_zip> --port <path_to_port_zip> --pac
 | `--skip-preflight` | Skip the preflight phase (not recommended) | `false` |
 | `--preflight-strict` | Treat risk findings as failures (strict gating) | `false` |
 | `--preflight-report` | Output path for preflight JSON report | `build/preflight-report.json` |
+| `--enable-snapshots` | Capture workspace snapshots at key workflow stages | `false` |
+| `--snapshot-dir` | Snapshot directory (defaults to `<work-dir>/snapshots`) | `null` |
+| `--rollback-to-snapshot` | Restore target workspace from a named snapshot and exit | `null` |
 
 ---
 

src/app/cli.py

@@ -61,6 +61,21 @@ def build_parser() -> argparse.ArgumentParser:
         default="build/preflight-report.json",
         help="Path to write preflight JSON report (default: build/preflight-report.json)",
     )
+    parser.add_argument(
+        "--enable-snapshots",
+        action="store_true",
+        help="Capture workflow snapshots at key stages",
+    )
+    parser.add_argument(
+        "--snapshot-dir",
+        default=None,
+        help="Snapshot directory (default: <work-dir>/snapshots)",
+    )
+    parser.add_argument(
+        "--rollback-to-snapshot",
+        default=None,
+        help="Restore target workspace from the named snapshot and exit",
+    )
     parser.add_argument(
         "--phases",
         nargs="+",

src/app/snapshots.py

@@ -0,0 +1,93 @@
+"""Stage snapshot and rollback helpers for workflow checkpoints."""
+
+from __future__ import annotations
+
+import json
+import logging
+import shutil
+from dataclasses import dataclass
+from datetime import datetime, timezone
+from pathlib import Path
+from typing import Any, cast
+
+
+@dataclass
+class SnapshotRecord:
+    """Metadata for a captured stage snapshot."""
+
+    name: str
+    created_at: str
+    source: str
+    snapshot_path: str
+
+
+class StageSnapshotManager:
+    """Capture and restore workflow stage snapshots."""
+
+    def __init__(self, root_dir: str | Path, logger: logging.Logger) -> None:
+        self.root_dir = Path(root_dir).resolve()
+        self.logger = logger
+        self.root_dir.mkdir(parents=True, exist_ok=True)
+        self.index_path = self.root_dir / "index.json"
+
+    def _load_index(self) -> dict[str, Any]:
+        if not self.index_path.exists():
+            return {"snapshots": []}
+        try:
+            loaded = json.loads(self.index_path.read_text(encoding="utf-8"))
+            if isinstance(loaded, dict):
+                return cast(dict[str, Any], loaded)
+            self.logger.warning("Snapshot index root is not an object. Reinitializing index.")
+            return {"snapshots": []}
+        except json.JSONDecodeError:
+            self.logger.warning("Snapshot index is invalid JSON. Reinitializing index.")
+            return {"snapshots": []}
+
+    def _save_index(self, data: dict[str, Any]) -> None:
+        self.index_path.write_text(json.dumps(data, indent=2), encoding="utf-8")
+
+    def capture(self, name: str, source_dir: str | Path) -> Path:
+        """Capture a snapshot for a workflow stage."""
+        source = Path(source_dir).resolve()
+        if not source.exists():
+            raise FileNotFoundError(f"Snapshot source does not exist: {source}")
+
+        snapshot_dir = self.root_dir / name
+        if snapshot_dir.exists():
+            shutil.rmtree(snapshot_dir)
+        shutil.copytree(source, snapshot_dir, dirs_exist_ok=True)
+
+        index = self._load_index()
+        snapshots = [
+            entry for entry in index.get("snapshots", []) if entry.get("name") != name
+        ]
+        snapshots.append(
+            {
+                "name": name,
+                "created_at": datetime.now(timezone.utc).isoformat(),
+                "source": str(source),
+                "snapshot_path": str(snapshot_dir),
+            }
+        )
+        index["snapshots"] = snapshots
+        self._save_index(index)
+        self.logger.info(f"Snapshot captured: {name} -> {snapshot_dir}")
+        return snapshot_dir
+
+    def restore(self, name: str, target_dir: str | Path) -> Path:
+        """Restore target directory from a named snapshot."""
+        snapshot_dir = self.root_dir / name
+        if not snapshot_dir.exists():
+            raise FileNotFoundError(f"Snapshot not found: {name}")
+
+        target = Path(target_dir).resolve()
+        if target.exists():
+            shutil.rmtree(target)
+        shutil.copytree(snapshot_dir, target, dirs_exist_ok=True)
+        self.logger.info(f"Snapshot restored: {name} -> {target}")
+        return target
+
+    def list_snapshot_names(self) -> list[str]:
+        """Return available snapshot names from index."""
+        index = self._load_index()
+        return [entry["name"] for entry in index.get("snapshots", []) if "name" in entry]

src/app/workflow.py

@@ -7,6 +7,7 @@
 
 from src.app.bootstrap import clean_work_dir, initialize_cache_manager
 from src.app.preflight import run_preflight, save_preflight_report
+from src.app.snapshots import StageSnapshotManager
 from src.core.config_loader import load_device_config
 from src.core.context import PortingContext
 from src.core.modifiers import FirmwareModifier, FrameworkModifier, RomModifier, UnifiedModifier
@@ -148,6 +149,28 @@ def execute_porting(args, logger: logging.Logger) -> int:
 
     resolve_remote_inputs(args, is_official_modify, logger)
 
+    work_dir, stock_work_dir, port_work_dir, target_work_dir = resolve_work_paths(args.work_dir)
+    snapshot_manager = (
+        StageSnapshotManager(args.snapshot_dir or (work_dir / "snapshots"), logger)
+        if args.enable_snapshots or args.rollback_to_snapshot
+        else None
+    )
+
+    if args.rollback_to_snapshot:
+        if not snapshot_manager:
+            logger.error("Snapshot manager is not available.")
+            return 1
+        try:
+            snapshot_manager.restore(args.rollback_to_snapshot, target_work_dir)
+            logger.info(f"Rollback completed from snapshot: {args.rollback_to_snapshot}")
+            return 0
+        except FileNotFoundError as exc:
+            logger.error(str(exc))
+            available = snapshot_manager.list_snapshot_names()
+            if available:
+                logger.info(f"Available snapshots: {', '.join(available)}")
+            return 2
+
     if not args.skip_preflight:
         preflight_report = run_preflight(args, is_official_modify, logger)
         report_path = save_preflight_report(preflight_report, args.preflight_report)
@@ -163,8 +186,6 @@ def execute_porting(args, logger: logging.Logger) -> int:
         logger.warning("Ignoring --preflight-only because --skip-preflight is set.")
         return 0
 
-    work_dir, stock_work_dir, port_work_dir, target_work_dir = resolve_work_paths(args.work_dir)
-
     if args.clean:
         clean_work_dir(work_dir, logger)
 
@@ -183,6 +204,8 @@ def execute_porting(args, logger: logging.Logger) -> int:
     ctx.cache_manager = cache_manager
     ctx.eu_bundle = args.eu_bundle
     ctx.initialize_target(clean_existing=True)
+    if snapshot_manager:
+        snapshot_manager.capture("phase2_initialized", target_work_dir)
 
     stock_device_code = (
         stock.get_prop("ro.product.name_for_attestation")
@@ -205,7 +228,12 @@ def execute_porting(args, logger: logging.Logger) -> int:
 
     phases_to_run = args.phases if args.phases else list(DEFAULT_PHASES)
     run_modification_phases(ctx, phases_to_run, logger)
+    if snapshot_manager:
+        snapshot_manager.capture("phase3_modified", target_work_dir)
+
     run_repacking(ctx, phases_to_run, pack_type, fs_type, target_work_dir, logger)
+    if snapshot_manager and ("repack" in phases_to_run or phases_to_run == DEFAULT_PHASES):
+        snapshot_manager.capture("phase4_repacked", target_work_dir)
 
     logger.info("=" * 70)
     logger.info("Porting completed successfully!")

tests/test_cli.py

@@ -29,3 +29,21 @@ def test_parse_args_accepts_preflight_flags():
     assert args.skip_preflight is False
     assert args.preflight_strict is False
     assert args.preflight_report == "out/preflight.json"
+
+
+def test_parse_args_accepts_snapshot_flags():
+    args = parse_args(
+        [
+            "--stock",
+            "stock.zip",
+            "--enable-snapshots",
+            "--snapshot-dir",
+            "build/snapshots",
+            "--rollback-to-snapshot",
+            "phase3_modified",
+        ]
+    )
+
+    assert args.enable_snapshots is True
+    assert args.snapshot_dir == "build/snapshots"
+    assert args.rollback_to_snapshot == "phase3_modified"

tests/test_snapshots.py

@@ -0,0 +1,30 @@
+import logging
+from pathlib import Path
+
+from src.app.snapshots import StageSnapshotManager
+
+
+def test_snapshot_capture_and_restore(tmp_path: Path):
+    source = tmp_path / "target"
+    source.mkdir(parents=True)
+    (source / "marker.txt").write_text("v1", encoding="utf-8")
+
+    manager = StageSnapshotManager(tmp_path / "snapshots", logging.getLogger("test"))
+    manager.capture("phase2_initialized", source)
+
+    (source / "marker.txt").write_text("v2", encoding="utf-8")
+    manager.restore("phase2_initialized", source)
+
+    assert (source / "marker.txt").read_text(encoding="utf-8") == "v1"
+
+
+def test_snapshot_list_names(tmp_path: Path):
+    source = tmp_path / "target"
+    source.mkdir(parents=True)
+    (source / "a.txt").write_text("ok", encoding="utf-8")
+
+    manager = StageSnapshotManager(tmp_path / "snapshots", logging.getLogger("test"))
+    manager.capture("phase2_initialized", source)
+    manager.capture("phase3_modified", source)
+
+    assert manager.list_snapshot_names() == ["phase2_initialized", "phase3_modified"]

tests/test_workflow.py

@@ -26,6 +26,9 @@ def make_args(**overrides):
         "skip_preflight": False,
         "preflight_strict": False,
         "preflight_report": "build/preflight-report.json",
+        "enable_snapshots": False,
+        "snapshot_dir": None,
+        "rollback_to_snapshot": None,
     }
     base.update(overrides)
     return Namespace(**base)
@@ -174,3 +177,75 @@ def test_execute_porting_strict_preflight_treats_risks_as_failures():
 
         assert execute_porting(args, logger) == 2
         run_preflight_mock.return_value.has_failures.assert_called_once_with(strict=True)
+
+
+def test_execute_porting_restores_snapshot_and_exits():
+    logger = MagicMock()
+    args = make_args(rollback_to_snapshot="phase3_modified")
+
+    with (
+        patch("src.app.workflow.initialize_cache_manager") as bootstrap,
+        patch("src.app.workflow.log_run_configuration"),
+        patch("src.app.workflow.OtaToolsManager") as otatools_manager_cls,
+        patch("src.app.workflow.resolve_remote_inputs"),
+        patch("src.app.workflow.StageSnapshotManager") as snapshot_manager_cls,
+        patch("src.app.workflow.resolve_work_paths") as resolve_work_paths,
+    ):
+        bootstrap.return_value.exit_code = None
+        bootstrap.return_value.cache_manager = None
+        otatools_manager_cls.return_value.ensure_otatools.return_value = True
+        resolve_work_paths.return_value = (
+            MagicMock(),
+            MagicMock(),
+            MagicMock(),
+            MagicMock(),
+        )
+
+        assert execute_porting(args, logger) == 0
+
+    snapshot_manager_cls.return_value.restore.assert_called_once()
+
+
+def test_execute_porting_captures_snapshots_when_enabled():
+    logger = MagicMock()
+    args = make_args(enable_snapshots=True)
+
+    with (
+        patch("src.app.workflow.initialize_cache_manager") as bootstrap,
+        patch("src.app.workflow.log_run_configuration"),
+        patch("src.app.workflow.OtaToolsManager") as otatools_manager_cls,
+        patch("src.app.workflow.resolve_remote_inputs"),
+        patch("src.app.workflow.run_preflight") as run_preflight_mock,
+        patch("src.app.workflow.save_preflight_report"),
+        patch("src.app.workflow.StageSnapshotManager") as snapshot_manager_cls,
+        patch("src.app.workflow.resolve_work_paths") as resolve_work_paths,
+        patch("src.app.workflow.RomPackage") as rom_package_cls,
+        patch("src.app.workflow.PortingContext") as porting_context_cls,
+        patch("src.app.workflow.load_device_config", return_value={}),
+        patch("src.app.workflow.determine_pack_settings", return_value=("payload", "erofs")),
+        patch("src.app.workflow.run_modification_phases"),
+        patch("src.app.workflow.run_repacking"),
+    ):
+        bootstrap.return_value.exit_code = None
+        bootstrap.return_value.cache_manager = None
+        otatools_manager_cls.return_value.ensure_otatools.return_value = True
+        run_preflight_mock.return_value.has_failures.return_value = False
+        resolve_work_paths.return_value = (
+            MagicMock(),
+            MagicMock(),
+            MagicMock(),
+            MagicMock(),
+        )
+        stock = rom_package_cls.return_value
+        porting_context = porting_context_cls.return_value
+        porting_context.stock = stock
+        porting_context.device_config = {}
+
+        assert execute_porting(args, logger) == 0
+
+    capture_calls = snapshot_manager_cls.return_value.capture.call_args_list
+    assert [call.args[0] for call in capture_calls] == [
+        "phase2_initialized",
+        "phase3_modified",
+        "phase4_repacked",
+    ]