"""Health state and quota tracking per profile and model family. Features: - Atomic file write + temporary file replace for router_state.json. - Thread-safe in-memory cache and automatic cooldown expiration. - Model family vs profile-level status tracking. """ from __future__ import annotations import json import os import tempfile import threading import time from dataclasses import asdict, dataclass, field from pathlib import Path from typing import Any, Dict, Optional from antigravity_provider import paths HEALTHY = "healthy" IN_USE = "in-use" QUOTA_EXHAUSTED = "quota-exhausted" RATE_LIMITED = "rate-limited" COOLDOWN = "cooldown" AUTH_REQUIRED = "auth-required" DISABLED = "disabled" UNHEALTHY = "unhealthy" @dataclass class FamilyHealthRecord: family: str state: str = HEALTHY reset_at: Optional[float] = None reason: Optional[str] = None last_error: Optional[str] = None error_count: int = 0 success_count: int = 0 simulated: bool = False @dataclass class ProfileHealthRecord: profile_id: str overall_state: str = HEALTHY families: dict[str, FamilyHealthRecord] = field(default_factory=dict) active_leases: int = 0 last_used: Optional[float] = None last_success: Optional[float] = None last_error: Optional[str] = None simulated: bool = False def extract_model_family(model_name: Optional[str]) -> str: """Extract model family prefix (e.g. gemini, claude, gpt, deepseek, kimi, qwen, grok, glm).""" if not model_name: return "default" m = model_name.lower().replace("google-antigravity/", "").replace("openai/", "").replace("moonshotai/", "") for family in ("gemini", "claude", "gpt", "o3", "o1", "deepseek", "kimi", "qwen", "grok", "glm", "mimo", "minimax"): if family in m: return family return "default" class _FileLock: """Interprocess file lock supporting Windows (msvcrt) and Unix (fcntl).""" def __init__(self, lock_path: Path): self.lock_path = lock_path self._fd: Optional[int] = None def __enter__(self): try: self.lock_path.parent.mkdir(parents=True, exist_ok=True) self._fd = os.open(str(self.lock_path), os.O_CREAT | os.O_RDWR) if os.name == "nt": import msvcrt msvcrt.locking(self._fd, msvcrt.LK_NBLCK, 1) else: import fcntl fcntl.flock(self._fd, fcntl.LOCK_EX | fcntl.LOCK_NB) except Exception: pass return self def __exit__(self, exc_type, exc_val, exc_tb): if self._fd is not None: try: if os.name == "nt": import msvcrt try: msvcrt.locking(self._fd, msvcrt.LK_UNLCK, 1) except Exception: pass else: import fcntl try: fcntl.flock(self._fd, fcntl.LOCK_UN) except Exception: pass os.close(self._fd) except Exception: pass self._fd = None class HealthTracker: """Thread-safe health tracker for router profiles and model families with atomic disk persistence and interprocess locking.""" def __init__(self, state_file: Optional[Path] = None): if state_file is None: state_file = paths.get_router_state_path() self.state_file = state_file self.lock_file = self.state_file.with_suffix(".lock") self._lock = threading.RLock() self._profiles: dict[str, ProfileHealthRecord] = {} self._load_state() def _load_state(self) -> None: if not self.state_file.is_file(): return try: raw = json.loads(self.state_file.read_text(encoding="utf-8")) for pid, pdata in raw.get("profiles", {}).items(): record = ProfileHealthRecord( profile_id=pid, overall_state=pdata.get("overall_state", HEALTHY), last_used=pdata.get("last_used"), last_success=pdata.get("last_success"), last_error=pdata.get("last_error"), simulated=pdata.get("simulated", False), ) for fname, fdata in pdata.get("families", {}).items(): record.families[fname] = FamilyHealthRecord( family=fname, state=fdata.get("state", HEALTHY), reset_at=fdata.get("reset_at"), reason=fdata.get("reason"), last_error=fdata.get("last_error"), error_count=fdata.get("error_count", 0), success_count=fdata.get("success_count", 0), simulated=fdata.get("simulated", False), ) self._profiles[pid] = record except Exception: pass def _save_state(self) -> None: """Atomically persist health state to disk with interprocess locking and temporary file replace.""" try: self.state_file.parent.mkdir(parents=True, exist_ok=True) data: dict[str, Any] = {"profiles": {}} for pid, precord in self._profiles.items(): pdict = { "overall_state": precord.overall_state, "last_used": precord.last_used, "last_success": precord.last_success, "last_error": precord.last_error, "simulated": precord.simulated, "families": {}, } for fname, frecord in precord.families.items(): pdict["families"][fname] = { "state": frecord.state, "reset_at": frecord.reset_at, "reason": frecord.reason, "last_error": frecord.last_error, "error_count": frecord.error_count, "success_count": frecord.success_count, "simulated": frecord.simulated, } data["profiles"][pid] = pdict serialized = json.dumps(data, indent=2, ensure_ascii=False) + "\n" with _FileLock(self.lock_file): # Atomic file replace tmp_fd, tmp_path = tempfile.mkstemp( dir=str(self.state_file.parent), prefix="router_state_", suffix=".tmp", ) with os.fdopen(tmp_fd, "w", encoding="utf-8") as f: f.write(serialized) # Atomic replace (works on Windows & POSIX in Python 3.3+) os.replace(tmp_path, str(self.state_file)) except Exception: pass def get_or_create(self, profile_id: str) -> ProfileHealthRecord: with self._lock: if profile_id not in self._profiles: self._profiles[profile_id] = ProfileHealthRecord(profile_id=profile_id) rec = self._profiles[profile_id] try: from .session_affinity import LeaseManager rec.active_leases = LeaseManager.get().active_count(profile_id) except Exception: pass return rec def is_healthy(self, profile_id: str, model_name: Optional[str] = None) -> bool: """Check if profile (and specified model family) is healthy and ready for requests.""" with self._lock: record = self.get_or_create(profile_id) now = time.time() if record.overall_state == DISABLED: return False # Check profile-level default family if "default" in record.families: def_rec = record.families["default"] if def_rec.state in (QUOTA_EXHAUSTED, RATE_LIMITED, COOLDOWN): if def_rec.reset_at and now >= def_rec.reset_at: def_rec.state = HEALTHY def_rec.reset_at = None def_rec.simulated = False else: return False family = extract_model_family(model_name) if family in record.families: frec = record.families[family] if frec.state == QUOTA_EXHAUSTED: if frec.reset_at and now >= frec.reset_at: # Expired cooldown -> reset to healthy frec.state = HEALTHY frec.reset_at = None frec.simulated = False self._save_state() return True return False if frec.state in (RATE_LIMITED, COOLDOWN): if frec.reset_at and now >= frec.reset_at: frec.state = HEALTHY frec.reset_at = None frec.simulated = False self._save_state() return True return False if frec.state in (AUTH_REQUIRED, UNHEALTHY, DISABLED): return False if record.overall_state in (AUTH_REQUIRED, UNHEALTHY, DISABLED, QUOTA_EXHAUSTED, RATE_LIMITED): return False return True def mark_success(self, profile_id: str, model_name: Optional[str] = None) -> None: with self._lock: record = self.get_or_create(profile_id) now = time.time() record.last_used = now record.last_success = now record.overall_state = HEALTHY record.simulated = False family = extract_model_family(model_name) if family in record.families: frec = record.families[family] frec.state = HEALTHY frec.reset_at = None frec.success_count += 1 frec.simulated = False if "default" in record.families: record.families["default"].state = HEALTHY record.families["default"].reset_at = None record.families["default"].simulated = False self._save_state() def mark_quota_exhausted( self, profile_id: str, model_name: Optional[str] = None, duration: int = 1800, reason: Optional[str] = None, simulated: bool = False, ) -> None: with self._lock: record = self.get_or_create(profile_id) now = time.time() record.last_used = now record.last_error = reason record.simulated = simulated if not model_name or model_name == "default": record.overall_state = QUOTA_EXHAUSTED family = extract_model_family(model_name) if family not in record.families: record.families[family] = FamilyHealthRecord(family=family) frec = record.families[family] frec.state = QUOTA_EXHAUSTED frec.reset_at = now + duration frec.reason = reason frec.last_error = reason frec.error_count += 1 frec.simulated = simulated self._save_state() def simulate_quota( self, profile_id: str, duration: int = 1800, model_family: Optional[str] = None, ) -> None: """Simulate quota exhaustion on a profile for testing.""" self.mark_quota_exhausted( profile_id=profile_id, model_name=model_family, duration=duration, reason="Simulated Quota Exhaustion", simulated=True, ) def mark_rate_limited( self, profile_id: str, model_name: Optional[str] = None, duration: int = 60, reason: Optional[str] = None, ) -> None: with self._lock: record = self.get_or_create(profile_id) now = time.time() record.last_used = now record.last_error = reason family = extract_model_family(model_name) if family not in record.families: record.families[family] = FamilyHealthRecord(family=family) frec = record.families[family] frec.state = RATE_LIMITED frec.reset_at = now + duration frec.reason = reason frec.last_error = reason frec.error_count += 1 self._save_state() def mark_auth_required(self, profile_id: str, reason: Optional[str] = None) -> None: with self._lock: record = self.get_or_create(profile_id) record.overall_state = AUTH_REQUIRED record.last_error = reason self._save_state() def clear_cooldown(self, profile_id: Optional[str] = None, model_name: Optional[str] = None) -> None: with self._lock: if profile_id is None: for rec in self._profiles.values(): rec.overall_state = HEALTHY rec.simulated = False for frec in rec.families.values(): frec.state = HEALTHY frec.reset_at = None frec.simulated = False self._save_state() return if profile_id not in self._profiles: return record = self._profiles[profile_id] record.overall_state = HEALTHY record.simulated = False if model_name: family = extract_model_family(model_name) if family in record.families: record.families[family].state = HEALTHY record.families[family].reset_at = None record.families[family].simulated = False else: for frec in record.families.values(): frec.state = HEALTHY frec.reset_at = None frec.simulated = False self._save_state()