1"""

2 Version Rollback — Safe rollback to any previous version.

3

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

5 and provides CLI for instant rollback.

6

7 Usage:

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

13 from __future__ import annotations

14

15 import hashlib

16 import json

17 import subprocess

18 import sys

19 from dataclasses import dataclass, field

20 from datetime import datetime, timezone

21 from pathlib import Path

22 from typing import Optional

23

24

25# ── Models ──

26

27@dataclass

28 class 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

36

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 }

46

47 @classmethod

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

49 return cls(**d)

50

51

52# ── Rollback Manager ──

53

54 class RollbackManager:

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

56

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

63

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)

70

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

72

73 # ── Archive ──

74

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

80

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

86

87 # Copy to archive

88 dest = self._wheels_dir / filename

89 import shutil

90 shutil.copy2(src, dest)

91

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

95

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 )

104

105 # Update history

106 self._history.append(entry)

107 self._save_history()

108

109 return entry

110

111 # ── Rollback ──

112

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

114 """Rollback to a previously archived version.

115

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

121

122 Args:

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

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

125

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

136

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

141

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

147

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

152

153 current_version = self._get_installed_version()

154

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

159

160 if dry_run:

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

162 return True

163

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 )

169

170 if result.returncode != 0:

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

172 return False

173

174 # Mark current as inactive

175 for entry in self._history:

176 if entry.version == current_version:

177 entry.active = False

178

179 # Mark target as active

180 target.active = True

181 self._save_history()

182

183 new_version = self._get_installed_version()

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

185 return True

186

187 # ── List ──

188

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

203

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 []

216

217 # ── Verify ──

218

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 )

226

227 pyapi_versions = self.list_pypi_versions()

228

229 for entry in entries:

230 wheel = Path(entry.wheel_path)

231 issues = []

232

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

239

240 if entry.version not in pyapi_versions:

241 issues.append("not on PyPI")

242

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 }

249

250 return results

251

252 # ── Clean ──

253

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

266

267 # ── Internal ──

268

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 []

277

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 )

282

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

290

291 @staticmethod

292 def _get_installed_version() -> str:

293 try:

294 import agentos

295 return agentos.__version__

296 except Exception:

297 return "unknown"

298

299

300# ── CLI Entry ──

301

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

303 """CLI entry for agentos rollback command.

304

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()

314

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

328

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

338

339 if "--prune" in args:

340 removed = mgr.prune()

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

342 return 0

343

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

353

354 # Default: rollback

355 if not args:

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

357 return 1

358

359 target = args[0]

360 success = mgr.rollback(target)

361 return 0 if success else 1