Coverage for agentos/cli/rollback.py: 0%

1""" 

2Version Rollback — Safe rollback to any previous version. 

3 

4Keeps a local archive of all pushed wheels (.agentos/wheels/) 

5and provides CLI for instant rollback. 

6 

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""" 

12 

13from __future__ import annotations 

14 

15import hashlib 

16import json 

17import subprocess 

18import sys 

19from dataclasses import dataclass 

20from datetime import datetime, timezone 

21from pathlib import Path 

22 

23 

24# ── Models ── 

25 

26@dataclass 

27class VersionEntry: 

28 """Record of a pushed version.""" 

29 version: str 

30 pushed_at: str 

31 wheel_path: str 

32 wheel_size: int 

33 sha256: str 

34 active: bool = True # False if rolled back from 

35 

36 def to_dict(self) -> dict: 

37 return { 

38 "version": self.version, 

39 "pushed_at": self.pushed_at, 

40 "wheel_path": self.wheel_path, 

41 "wheel_size": self.wheel_size, 

42 "sha256": self.sha256, 

43 "active": self.active, 

44 } 

45 

46 @classmethod 

47 def from_dict(cls, d: dict) -> "VersionEntry": 

48 return cls(**d) 

49 

50 

51# ── Rollback Manager ── 

52 

53class RollbackManager: 

54 """Safe version rollback with local wheel archive. 

55 

56 Archive: ~/.agentos/rollback/ 

57 ├── history.json # Version records 

58 └── wheels/ # Archived .whl files 

59 ├── nexus_agentos-1.7.4-py3-none-any.whl 

60 └── nexus_agentos-1.7.5-py3-none-any.whl 

61 """ 

62 

63 def __init__(self, archive_dir: str = ""): 

64 self._root = Path(archive_dir) if archive_dir else Path.home() / ".agentos" / "rollback" 

65 self._history_path = self._root / "history.json" 

66 self._wheels_dir = self._root / "wheels" 

67 self._root.mkdir(parents=True, exist_ok=True) 

68 self._wheels_dir.mkdir(parents=True, exist_ok=True) 

69 

70 self._history: list[VersionEntry] = self._load_history() 

71 

72 # ── Archive ── 

73 

74 def archive(self, wheel_path: str | Path) -> VersionEntry: 

75 """Archive a wheel after pushing to PyPI. Call after twine upload.""" 

76 src = Path(wheel_path) 

77 if not src.exists(): 

78 raise FileNotFoundError(f"Wheel not found: {src}") 

79 

80 # Parse version from filename 

81 filename = src.name 

82 version = self._parse_version(filename) 

83 if not version: 

84 raise ValueError(f"Cannot parse version from {filename}") 

85 

86 # Copy to archive 

87 dest = self._wheels_dir / filename 

88 import shutil 

89 shutil.copy2(src, dest) 

90 

91 # Compute hash 

92 sha = hashlib.sha256(dest.read_bytes()).hexdigest() 

93 pushed_at = datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ") 

94 

95 entry = VersionEntry( 

96 version=version, 

97 pushed_at=pushed_at, 

98 wheel_path=str(dest), 

99 wheel_size=dest.stat().st_size, 

100 sha256=sha, 

101 active=True, 

102 ) 

103 

104 # Update history 

105 self._history.append(entry) 

106 self._save_history() 

107 

108 return entry 

109 

110 # ── Rollback ── 

111 

112 def rollback(self, target_version: str, dry_run: bool = False) -> bool: 

113 """Rollback to a previously archived version. 

114 

115 Steps: 

116 1. Find the target version's wheel in archive 

117 2. Verify SHA256 integrity 

118 3. pip install the archived wheel 

119 4. Mark current as inactive, target as active 

120 

121 Args: 

122 target_version: e.g. '1.7.3' or '1.7.4' 

123 dry_run: If True, validate only, don't install. 

124 

125 Returns: 

126 True if rollback succeeded (or would succeed in dry_run). 

127 """ 

128 # Find target 

129 target = None 

130 for entry in self._history: 

131 if entry.version == target_version and entry.active is False: 

132 target = entry 

133 elif entry.version == target_version and Path(entry.wheel_path).exists(): 

134 target = entry 

135 

136 if not target: 

137 available = [e.version for e in self._history if Path(e.wheel_path).exists()] 

138 print(f"Version {target_version} not found in archive. Available: {available}") 

139 return False 

140 

141 # Verify integrity 

142 wheel = Path(target.wheel_path) 

143 if not wheel.exists(): 

144 print(f"Wheel file missing: {target.wheel_path}") 

145 return False 

146 

147 actual_sha = hashlib.sha256(wheel.read_bytes()).hexdigest() 

148 if actual_sha != target.sha256: 

149 print(f"SHA256 mismatch! Expected: {target.sha256[:16]}..., Got: {actual_sha[:16]}...") 

150 return False 

151 

152 current_version = self._get_installed_version() 

153 

154 print(f"Rollback: {current_version} → {target_version}") 

155 print(f" Wheel: {wheel}") 

156 print(f" SHA256: {target.sha256[:16]}... (verified)") 

157 print(f" Size: {target.wheel_size / 1024:.0f} KB") 

158 

159 if dry_run: 

160 print(" [DRY RUN — no changes made]") 

161 return True 

162 

163 # pip install 

164 result = subprocess.run( 

165 [sys.executable, "-m", "pip", "install", "--force-reinstall", "--no-deps", str(wheel)], 

166 capture_output=True, text=True, 

167 ) 

168 

169 if result.returncode != 0: 

170 print(f"pip install failed:\n{result.stderr}") 

171 return False 

172 

173 # Mark current as inactive 

174 for entry in self._history: 

175 if entry.version == current_version: 

176 entry.active = False 

177 

178 # Mark target as active 

179 target.active = True 

180 self._save_history() 

181 

182 new_version = self._get_installed_version() 

183 print(f"Rollback successful: {current_version} → {new_version}") 

184 return True 

185 

186 # ── List ── 

187 

188 def list_versions(self) -> list[dict]: 

189 """List all archived versions with status.""" 

190 installed = self._get_installed_version() 

191 result = [] 

192 for entry in sorted(self._history, key=lambda e: e.pushed_at, reverse=True): 

193 result.append({ 

194 "version": entry.version, 

195 "pushed_at": entry.pushed_at, 

196 "wheel_size_kb": entry.wheel_size // 1024, 

197 "active": entry.active, 

198 "current": entry.version == installed, 

199 "wheel_exists": Path(entry.wheel_path).exists(), 

200 }) 

201 return result 

202 

203 def list_pypi_versions(self) -> list[str]: 

204 """List all versions available on PyPI.""" 

205 import urllib.request 

206 import json 

207 try: 

208 url = "https://pypi.org/pypi/nexus-agentos/json" 

209 req = urllib.request.Request(url, headers={"Accept": "application/json"}) 

210 with urllib.request.urlopen(req, timeout=10) as resp: 

211 data = json.loads(resp.read().decode()) 

212 return sorted(data.get("releases", {}).keys(), reverse=True) 

213 except Exception: 

214 return [] 

215 

216 # ── Verify ── 

217 

218 def verify(self, version: str = "") -> dict: 

219 """Verify a specific version's wheel integrity, or all.""" 

220 results = {} 

221 entries = ( 

222 [e for e in self._history if e.version == version] 

223 if version else self._history 

224 ) 

225 

226 pyapi_versions = self.list_pypi_versions() 

227 

228 for entry in entries: 

229 wheel = Path(entry.wheel_path) 

230 issues = [] 

231 

232 if not wheel.exists(): 

233 issues.append("wheel file missing") 

234 else: 

235 actual_sha = hashlib.sha256(wheel.read_bytes()).hexdigest() 

236 if actual_sha != entry.sha256: 

237 issues.append("SHA256 mismatch") 

238 

239 if entry.version not in pyapi_versions: 

240 issues.append("not on PyPI") 

241 

242 results[entry.version] = { 

243 "sha256_ok": not any("sha256" in i.lower() for i in issues), 

244 "wheel_exists": wheel.exists(), 

245 "on_pypi": entry.version in pyapi_versions, 

246 "issues": issues, 

247 } 

248 

249 return results 

250 

251 # ── Clean ── 

252 

253 def prune(self, keep_versions: int = 5) -> list[str]: 

254 """Remove old wheel files, keeping the N most recent.""" 

255 removed = [] 

256 entries = sorted(self._history, key=lambda e: e.pushed_at, reverse=True) 

257 for entry in entries[keep_versions:]: 

258 wheel = Path(entry.wheel_path) 

259 if wheel.exists(): 

260 wheel.unlink() 

261 removed.append(entry.version) 

262 self._history = entries[:keep_versions] 

263 self._save_history() 

264 return removed 

265 

266 # ── Internal ── 

267 

268 def _load_history(self) -> list[VersionEntry]: 

269 if self._history_path.exists(): 

270 try: 

271 data = json.loads(self._history_path.read_text()) 

272 return [VersionEntry.from_dict(d) for d in data] 

273 except (json.JSONDecodeError, KeyError): 

274 pass 

275 return [] 

276 

277 def _save_history(self) -> None: 

278 self._history_path.write_text( 

279 json.dumps([e.to_dict() for e in self._history], indent=2, ensure_ascii=False) 

280 ) 

281 

282 @staticmethod 

283 def _parse_version(filename: str) -> str: 

284 """Extract version from wheel filename: nexus_agentos-1.7.5-py3-none-any.whl → 1.7.5""" 

285 parts = filename.replace(".whl", "").split("-") 

286 if len(parts) >= 2: 

287 return parts[1].replace(".post", ".") # Normalize .postN suffix 

288 return "" 

289 

290 @staticmethod 

291 def _get_installed_version() -> str: 

292 try: 

293 import agentos 

294 return agentos.__version__ 

295 except Exception: 

296 return "unknown" 

297 

298 

299# ── CLI Entry ── 

300 

301def rollback_cli(args: list[str]) -> int: 

302 """CLI entry for agentos rollback command. 

303 

304 Usage: 

305 agentos rollback 1.7.4 # Rollback 

306 agentos rollback --list # List versions 

307 agentos rollback --verify # Verify all 

308 agentos rollback --verify 1.7.5 # Verify one 

309 agentos rollback --prune # Keep last 5 versions 

310 agentos rollback --archive dist/nexus_agentos-1.7.5-py3-none-any.whl 

311 """ 

312 mgr = RollbackManager() 

313 

314 if "--list" in args: 

315 versions = mgr.list_versions() 

316 if not versions: 

317 print("No versions in rollback archive. Archive a wheel first.") 

318 return 0 

319 print(f"{'Version':<12} {'Pushed':<22} {'Size':>8} {'Status':<10} {'On PyPI'}") 

320 print("-" * 70) 

321 for v in versions: 

322 status = "✓ current" if v["current"] else " active" if v["active"] else " inactive" 

323 size = f"{v['wheel_size_kb']} KB" 

324 pypi = "✓" if v.get("on_pypi", True) else "" 

325 print(f" {v['version']:<10} {v['pushed_at']:<22} {size:>8} {status:<10} {pypi:^6}") 

326 return 0 

327 

328 if "--verify" in args: 

329 idx = args.index("--verify") 

330 target = args[idx + 1] if idx + 1 < len(args) else "" 

331 results = mgr.verify(target) 

332 for ver, r in results.items(): 

333 ok = "✓" if not r["issues"] else "✗" 

334 issues = ", ".join(r["issues"]) if r["issues"] else "clean" 

335 print(f" {ok} {ver}: {issues}") 

336 return 0 if all(not r["issues"] for r in results.values()) else 1 

337 

338 if "--prune" in args: 

339 removed = mgr.prune() 

340 print(f"Pruned {len(removed)} old wheels: {removed}") 

341 return 0 

342 

343 if "--archive" in args: 

344 idx = args.index("--archive") 

345 wheel = args[idx + 1] if idx + 1 < len(args) else "" 

346 if not wheel: 

347 print("Usage: agentos rollback --archive <wheel_path>") 

348 return 1 

349 entry = mgr.archive(wheel) 

350 print(f"Archived: {entry.version} ({entry.wheel_size / 1024:.0f} KB)") 

351 return 0 

352 

353 # Default: rollback 

354 if not args: 

355 print("Usage: agentos rollback <version> [--list|--verify|--prune|--archive]") 

356 return 1 

357 

358 target = args[0] 

359 success = mgr.rollback(target) 

360 return 0 if success else 1