Coverage for agentos/cli/rollback.py: 0%
195 statements
« prev ^ index » next coverage.py v7.14.3, created at 2026-07-09 09:19 +0800
1"""
2Version Rollback — Safe rollback to any previous version.
4Keeps a local archive of all pushed wheels (.agentos/wheels/)
5and provides CLI for instant rollback.
7Usage:
8 agentos rollback 1.7.4 # Rollback to 1.7.4
9 agentos rollback --list # List available versions
10 agentos rollback --verify 1.7.5 # Verify a version's wheel integrity
11"""
13from __future__ import annotations
15import hashlib
16import json
17import subprocess
18import sys
19from dataclasses import dataclass
20from datetime import UTC, datetime
21from pathlib import Path
23# ── Models ──
26@dataclass
27class VersionEntry:
28 """Record of a pushed version."""
30 version: str
31 pushed_at: str
32 wheel_path: str
33 wheel_size: int
34 sha256: str
35 active: bool = True # False if rolled back from
37 def to_dict(self) -> dict:
38 return {
39 "version": self.version,
40 "pushed_at": self.pushed_at,
41 "wheel_path": self.wheel_path,
42 "wheel_size": self.wheel_size,
43 "sha256": self.sha256,
44 "active": self.active,
45 }
47 @classmethod
48 def from_dict(cls, d: dict) -> VersionEntry:
49 return cls(**d)
52# ── Rollback Manager ──
55class RollbackManager:
56 """Safe version rollback with local wheel archive.
58 Archive: ~/.agentos/rollback/
59 ├── history.json # Version records
60 └── wheels/ # Archived .whl files
61 ├── nexus_agentos-1.7.4-py3-none-any.whl
62 └── nexus_agentos-1.7.5-py3-none-any.whl
63 """
65 def __init__(self, archive_dir: str = ""):
66 self._root = Path(archive_dir) if archive_dir else Path.home() / ".agentos" / "rollback"
67 self._history_path = self._root / "history.json"
68 self._wheels_dir = self._root / "wheels"
69 self._root.mkdir(parents=True, exist_ok=True)
70 self._wheels_dir.mkdir(parents=True, exist_ok=True)
72 self._history: list[VersionEntry] = self._load_history()
74 # ── Archive ──
76 def archive(self, wheel_path: str | Path) -> VersionEntry:
77 """Archive a wheel after pushing to PyPI. Call after twine upload."""
78 src = Path(wheel_path)
79 if not src.exists():
80 raise FileNotFoundError(f"Wheel not found: {src}")
82 # Parse version from filename
83 filename = src.name
84 version = self._parse_version(filename)
85 if not version:
86 raise ValueError(f"Cannot parse version from {filename}")
88 # Copy to archive
89 dest = self._wheels_dir / filename
90 import shutil
92 shutil.copy2(src, dest)
94 # Compute hash
95 sha = hashlib.sha256(dest.read_bytes()).hexdigest()
96 pushed_at = datetime.now(UTC).strftime("%Y-%m-%dT%H:%M:%SZ")
98 entry = VersionEntry(
99 version=version,
100 pushed_at=pushed_at,
101 wheel_path=str(dest),
102 wheel_size=dest.stat().st_size,
103 sha256=sha,
104 active=True,
105 )
107 # Update history
108 self._history.append(entry)
109 self._save_history()
111 return entry
113 # ── Rollback ──
115 def rollback(self, target_version: str, dry_run: bool = False) -> bool:
116 """Rollback to a previously archived version.
118 Steps:
119 1. Find the target version's wheel in archive
120 2. Verify SHA256 integrity
121 3. pip install the archived wheel
122 4. Mark current as inactive, target as active
124 Args:
125 target_version: e.g. '1.7.3' or '1.7.4'
126 dry_run: If True, validate only, don't install.
128 Returns:
129 True if rollback succeeded (or would succeed in dry_run).
130 """
131 # Find target
132 target = None
133 for entry in self._history:
134 if entry.version == target_version and entry.active is False:
135 target = entry
136 elif entry.version == target_version and Path(entry.wheel_path).exists():
137 target = entry
139 if not target:
140 available = [e.version for e in self._history if Path(e.wheel_path).exists()]
141 print(f"Version {target_version} not found in archive. Available: {available}")
142 return False
144 # Verify integrity
145 wheel = Path(target.wheel_path)
146 if not wheel.exists():
147 print(f"Wheel file missing: {target.wheel_path}")
148 return False
150 actual_sha = hashlib.sha256(wheel.read_bytes()).hexdigest()
151 if actual_sha != target.sha256:
152 print(f"SHA256 mismatch! Expected: {target.sha256[:16]}..., Got: {actual_sha[:16]}...")
153 return False
155 current_version = self._get_installed_version()
157 print(f"Rollback: {current_version} → {target_version}")
158 print(f" Wheel: {wheel}")
159 print(f" SHA256: {target.sha256[:16]}... (verified)")
160 print(f" Size: {target.wheel_size / 1024:.0f} KB")
162 if dry_run:
163 print(" [DRY RUN — no changes made]")
164 return True
166 # pip install
167 result = subprocess.run(
168 [sys.executable, "-m", "pip", "install", "--force-reinstall", "--no-deps", str(wheel)],
169 capture_output=True,
170 text=True,
171 )
173 if result.returncode != 0:
174 print(f"pip install failed:\n{result.stderr}")
175 return False
177 # Mark current as inactive
178 for entry in self._history:
179 if entry.version == current_version:
180 entry.active = False
182 # Mark target as active
183 target.active = True
184 self._save_history()
186 new_version = self._get_installed_version()
187 print(f"Rollback successful: {current_version} → {new_version}")
188 return True
190 # ── List ──
192 def list_versions(self) -> list[dict]:
193 """List all archived versions with status."""
194 installed = self._get_installed_version()
195 result = []
196 for entry in sorted(self._history, key=lambda e: e.pushed_at, reverse=True):
197 result.append(
198 {
199 "version": entry.version,
200 "pushed_at": entry.pushed_at,
201 "wheel_size_kb": entry.wheel_size // 1024,
202 "active": entry.active,
203 "current": entry.version == installed,
204 "wheel_exists": Path(entry.wheel_path).exists(),
205 }
206 )
207 return result
209 def list_pypi_versions(self) -> list[str]:
210 """List all versions available on PyPI."""
211 import json
212 import urllib.request
214 try:
215 url = "https://pypi.org/pypi/nexus-agentos/json"
216 req = urllib.request.Request(url, headers={"Accept": "application/json"})
217 with urllib.request.urlopen(req, timeout=10) as resp:
218 data = json.loads(resp.read().decode())
219 return sorted(data.get("releases", {}).keys(), reverse=True)
220 except Exception:
221 return []
223 # ── Verify ──
225 def verify(self, version: str = "") -> dict:
226 """Verify a specific version's wheel integrity, or all."""
227 results = {}
228 entries = [e for e in self._history if e.version == version] if version else self._history
230 pyapi_versions = self.list_pypi_versions()
232 for entry in entries:
233 wheel = Path(entry.wheel_path)
234 issues = []
236 if not wheel.exists():
237 issues.append("wheel file missing")
238 else:
239 actual_sha = hashlib.sha256(wheel.read_bytes()).hexdigest()
240 if actual_sha != entry.sha256:
241 issues.append("SHA256 mismatch")
243 if entry.version not in pyapi_versions:
244 issues.append("not on PyPI")
246 results[entry.version] = {
247 "sha256_ok": not any("sha256" in i.lower() for i in issues),
248 "wheel_exists": wheel.exists(),
249 "on_pypi": entry.version in pyapi_versions,
250 "issues": issues,
251 }
253 return results
255 # ── Clean ──
257 def prune(self, keep_versions: int = 5) -> list[str]:
258 """Remove old wheel files, keeping the N most recent."""
259 removed = []
260 entries = sorted(self._history, key=lambda e: e.pushed_at, reverse=True)
261 for entry in entries[keep_versions:]:
262 wheel = Path(entry.wheel_path)
263 if wheel.exists():
264 wheel.unlink()
265 removed.append(entry.version)
266 self._history = entries[:keep_versions]
267 self._save_history()
268 return removed
270 # ── Internal ──
272 def _load_history(self) -> list[VersionEntry]:
273 if self._history_path.exists():
274 try:
275 data = json.loads(self._history_path.read_text())
276 return [VersionEntry.from_dict(d) for d in data]
277 except (json.JSONDecodeError, KeyError):
278 pass
279 return []
281 def _save_history(self) -> None:
282 self._history_path.write_text(
283 json.dumps([e.to_dict() for e in self._history], indent=2, ensure_ascii=False)
284 )
286 @staticmethod
287 def _parse_version(filename: str) -> str:
288 """Extract version from wheel filename: nexus_agentos-1.7.5-py3-none-any.whl → 1.7.5"""
289 parts = filename.replace(".whl", "").split("-")
290 if len(parts) >= 2:
291 return parts[1].replace(".post", ".") # Normalize .postN suffix
292 return ""
294 @staticmethod
295 def _get_installed_version() -> str:
296 try:
297 import agentos
299 return agentos.__version__
300 except Exception:
301 return "unknown"
304# ── CLI Entry ──
307def rollback_cli(args: list[str]) -> int:
308 """CLI entry for agentos rollback command.
310 Usage:
311 agentos rollback 1.7.4 # Rollback
312 agentos rollback --list # List versions
313 agentos rollback --verify # Verify all
314 agentos rollback --verify 1.7.5 # Verify one
315 agentos rollback --prune # Keep last 5 versions
316 agentos rollback --archive dist/nexus_agentos-1.7.5-py3-none-any.whl
317 """
318 mgr = RollbackManager()
320 if "--list" in args:
321 versions = mgr.list_versions()
322 if not versions:
323 print("No versions in rollback archive. Archive a wheel first.")
324 return 0
325 print(f"{'Version':<12} {'Pushed':<22} {'Size':>8} {'Status':<10} {'On PyPI'}")
326 print("-" * 70)
327 for v in versions:
328 status = "✓ current" if v["current"] else " active" if v["active"] else " inactive"
329 size = f"{v['wheel_size_kb']} KB"
330 pypi = "✓" if v.get("on_pypi", True) else ""
331 print(f" {v['version']:<10} {v['pushed_at']:<22} {size:>8} {status:<10} {pypi:^6}")
332 return 0
334 if "--verify" in args:
335 idx = args.index("--verify")
336 target = args[idx + 1] if idx + 1 < len(args) else ""
337 results = mgr.verify(target)
338 for ver, r in results.items():
339 ok = "✓" if not r["issues"] else "✗"
340 issues = ", ".join(r["issues"]) if r["issues"] else "clean"
341 print(f" {ok} {ver}: {issues}")
342 return 0 if all(not r["issues"] for r in results.values()) else 1
344 if "--prune" in args:
345 removed = mgr.prune()
346 print(f"Pruned {len(removed)} old wheels: {removed}")
347 return 0
349 if "--archive" in args:
350 idx = args.index("--archive")
351 wheel = args[idx + 1] if idx + 1 < len(args) else ""
352 if not wheel:
353 print("Usage: agentos rollback --archive <wheel_path>")
354 return 1
355 entry = mgr.archive(wheel)
356 print(f"Archived: {entry.version} ({entry.wheel_size / 1024:.0f} KB)")
357 return 0
359 # Default: rollback
360 if not args:
361 print("Usage: agentos rollback <version> [--list|--verify|--prune|--archive]")
362 return 1
364 target = args[0]
365 success = mgr.rollback(target)
366 return 0 if success else 1