Coverage for agentos/cli/rollback.py: 0%
196 statements
« prev ^ index » next coverage.py v7.14.3, created at 2026-07-02 09:59 +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, field
20from datetime import datetime, timezone
21from pathlib import Path
22from typing import Optional
25# ── Models ──
27@dataclass
28class VersionEntry:
29 """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 ──
54class RollbackManager:
55 """Safe version rollback with local wheel archive.
57 Archive: ~/.agentos/rollback/
58 ├── history.json # Version records
59 └── wheels/ # Archived .whl files
60 ├── nexus_agentos-1.7.4-py3-none-any.whl
61 └── nexus_agentos-1.7.5-py3-none-any.whl
62 """
64 def __init__(self, archive_dir: str = ""):
65 self._root = Path(archive_dir) if archive_dir else Path.home() / ".agentos" / "rollback"
66 self._history_path = self._root / "history.json"
67 self._wheels_dir = self._root / "wheels"
68 self._root.mkdir(parents=True, exist_ok=True)
69 self._wheels_dir.mkdir(parents=True, exist_ok=True)
71 self._history: list[VersionEntry] = self._load_history()
73 # ── Archive ──
75 def archive(self, wheel_path: str | Path) -> VersionEntry:
76 """Archive a wheel after pushing to PyPI. Call after twine upload."""
77 src = Path(wheel_path)
78 if not src.exists():
79 raise FileNotFoundError(f"Wheel not found: {src}")
81 # Parse version from filename
82 filename = src.name
83 version = self._parse_version(filename)
84 if not version:
85 raise ValueError(f"Cannot parse version from {filename}")
87 # Copy to archive
88 dest = self._wheels_dir / filename
89 import shutil
90 shutil.copy2(src, dest)
92 # Compute hash
93 sha = hashlib.sha256(dest.read_bytes()).hexdigest()
94 pushed_at = datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ")
96 entry = VersionEntry(
97 version=version,
98 pushed_at=pushed_at,
99 wheel_path=str(dest),
100 wheel_size=dest.stat().st_size,
101 sha256=sha,
102 active=True,
103 )
105 # Update history
106 self._history.append(entry)
107 self._save_history()
109 return entry
111 # ── Rollback ──
113 def rollback(self, target_version: str, dry_run: bool = False) -> bool:
114 """Rollback to a previously archived version.
116 Steps:
117 1. Find the target version's wheel in archive
118 2. Verify SHA256 integrity
119 3. pip install the archived wheel
120 4. Mark current as inactive, target as active
122 Args:
123 target_version: e.g. '1.7.3' or '1.7.4'
124 dry_run: If True, validate only, don't install.
126 Returns:
127 True if rollback succeeded (or would succeed in dry_run).
128 """
129 # Find target
130 target = None
131 for entry in self._history:
132 if entry.version == target_version and entry.active is False:
133 target = entry
134 elif entry.version == target_version and Path(entry.wheel_path).exists():
135 target = entry
137 if not target:
138 available = [e.version for e in self._history if Path(e.wheel_path).exists()]
139 print(f"Version {target_version} not found in archive. Available: {available}")
140 return False
142 # Verify integrity
143 wheel = Path(target.wheel_path)
144 if not wheel.exists():
145 print(f"Wheel file missing: {target.wheel_path}")
146 return False
148 actual_sha = hashlib.sha256(wheel.read_bytes()).hexdigest()
149 if actual_sha != target.sha256:
150 print(f"SHA256 mismatch! Expected: {target.sha256[:16]}..., Got: {actual_sha[:16]}...")
151 return False
153 current_version = self._get_installed_version()
155 print(f"Rollback: {current_version} → {target_version}")
156 print(f" Wheel: {wheel}")
157 print(f" SHA256: {target.sha256[:16]}... (verified)")
158 print(f" Size: {target.wheel_size / 1024:.0f} KB")
160 if dry_run:
161 print(f" [DRY RUN — no changes made]")
162 return True
164 # pip install
165 result = subprocess.run(
166 [sys.executable, "-m", "pip", "install", "--force-reinstall", "--no-deps", str(wheel)],
167 capture_output=True, text=True,
168 )
170 if result.returncode != 0:
171 print(f"pip install failed:\n{result.stderr}")
172 return False
174 # Mark current as inactive
175 for entry in self._history:
176 if entry.version == current_version:
177 entry.active = False
179 # Mark target as active
180 target.active = True
181 self._save_history()
183 new_version = self._get_installed_version()
184 print(f"Rollback successful: {current_version} → {new_version}")
185 return True
187 # ── List ──
189 def list_versions(self) -> list[dict]:
190 """List all archived versions with status."""
191 installed = self._get_installed_version()
192 result = []
193 for entry in sorted(self._history, key=lambda e: e.pushed_at, reverse=True):
194 result.append({
195 "version": entry.version,
196 "pushed_at": entry.pushed_at,
197 "wheel_size_kb": entry.wheel_size // 1024,
198 "active": entry.active,
199 "current": entry.version == installed,
200 "wheel_exists": Path(entry.wheel_path).exists(),
201 })
202 return result
204 def list_pypi_versions(self) -> list[str]:
205 """List all versions available on PyPI."""
206 import urllib.request
207 import json
208 try:
209 url = "https://pypi.org/pypi/nexus-agentos/json"
210 req = urllib.request.Request(url, headers={"Accept": "application/json"})
211 with urllib.request.urlopen(req, timeout=10) as resp:
212 data = json.loads(resp.read().decode())
213 return sorted(data.get("releases", {}).keys(), reverse=True)
214 except Exception:
215 return []
217 # ── Verify ──
219 def verify(self, version: str = "") -> dict:
220 """Verify a specific version's wheel integrity, or all."""
221 results = {}
222 entries = (
223 [e for e in self._history if e.version == version]
224 if version else self._history
225 )
227 pyapi_versions = self.list_pypi_versions()
229 for entry in entries:
230 wheel = Path(entry.wheel_path)
231 issues = []
233 if not wheel.exists():
234 issues.append("wheel file missing")
235 else:
236 actual_sha = hashlib.sha256(wheel.read_bytes()).hexdigest()
237 if actual_sha != entry.sha256:
238 issues.append(f"SHA256 mismatch")
240 if entry.version not in pyapi_versions:
241 issues.append("not on PyPI")
243 results[entry.version] = {
244 "sha256_ok": not any("sha256" in i.lower() for i in issues),
245 "wheel_exists": wheel.exists(),
246 "on_pypi": entry.version in pyapi_versions,
247 "issues": issues,
248 }
250 return results
252 # ── Clean ──
254 def prune(self, keep_versions: int = 5) -> list[str]:
255 """Remove old wheel files, keeping the N most recent."""
256 removed = []
257 entries = sorted(self._history, key=lambda e: e.pushed_at, reverse=True)
258 for entry in entries[keep_versions:]:
259 wheel = Path(entry.wheel_path)
260 if wheel.exists():
261 wheel.unlink()
262 removed.append(entry.version)
263 self._history = entries[:keep_versions]
264 self._save_history()
265 return removed
267 # ── Internal ──
269 def _load_history(self) -> list[VersionEntry]:
270 if self._history_path.exists():
271 try:
272 data = json.loads(self._history_path.read_text())
273 return [VersionEntry.from_dict(d) for d in data]
274 except (json.JSONDecodeError, KeyError):
275 pass
276 return []
278 def _save_history(self) -> None:
279 self._history_path.write_text(
280 json.dumps([e.to_dict() for e in self._history], indent=2, ensure_ascii=False)
281 )
283 @staticmethod
284 def _parse_version(filename: str) -> str:
285 """Extract version from wheel filename: nexus_agentos-1.7.5-py3-none-any.whl → 1.7.5"""
286 parts = filename.replace(".whl", "").split("-")
287 if len(parts) >= 2:
288 return parts[1].replace(".post", ".") # Normalize .postN suffix
289 return ""
291 @staticmethod
292 def _get_installed_version() -> str:
293 try:
294 import agentos
295 return agentos.__version__
296 except Exception:
297 return "unknown"
300# ── CLI Entry ──
302def rollback_cli(args: list[str]) -> int:
303 """CLI entry for agentos rollback command.
305 Usage:
306 agentos rollback 1.7.4 # Rollback
307 agentos rollback --list # List versions
308 agentos rollback --verify # Verify all
309 agentos rollback --verify 1.7.5 # Verify one
310 agentos rollback --prune # Keep last 5 versions
311 agentos rollback --archive dist/nexus_agentos-1.7.5-py3-none-any.whl
312 """
313 mgr = RollbackManager()
315 if "--list" in args:
316 versions = mgr.list_versions()
317 if not versions:
318 print("No versions in rollback archive. Archive a wheel first.")
319 return 0
320 print(f"{'Version':<12} {'Pushed':<22} {'Size':>8} {'Status':<10} {'On PyPI'}")
321 print("-" * 70)
322 for v in versions:
323 status = "✓ current" if v["current"] else " active" if v["active"] else " inactive"
324 size = f"{v['wheel_size_kb']} KB"
325 pypi = "✓" if v.get("on_pypi", True) else ""
326 print(f" {v['version']:<10} {v['pushed_at']:<22} {size:>8} {status:<10} {pypi:^6}")
327 return 0
329 if "--verify" in args:
330 idx = args.index("--verify")
331 target = args[idx + 1] if idx + 1 < len(args) else ""
332 results = mgr.verify(target)
333 for ver, r in results.items():
334 ok = "✓" if not r["issues"] else "✗"
335 issues = ", ".join(r["issues"]) if r["issues"] else "clean"
336 print(f" {ok} {ver}: {issues}")
337 return 0 if all(not r["issues"] for r in results.values()) else 1
339 if "--prune" in args:
340 removed = mgr.prune()
341 print(f"Pruned {len(removed)} old wheels: {removed}")
342 return 0
344 if "--archive" in args:
345 idx = args.index("--archive")
346 wheel = args[idx + 1] if idx + 1 < len(args) else ""
347 if not wheel:
348 print("Usage: agentos rollback --archive <wheel_path>")
349 return 1
350 entry = mgr.archive(wheel)
351 print(f"Archived: {entry.version} ({entry.wheel_size / 1024:.0f} KB)")
352 return 0
354 # Default: rollback
355 if not args:
356 print("Usage: agentos rollback <version> [--list|--verify|--prune|--archive]")
357 return 1
359 target = args[0]
360 success = mgr.rollback(target)
361 return 0 if success else 1