585 lines
23 KiB
Python
585 lines
23 KiB
Python
"""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
|
||
auth_error_at: Optional[float] = 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"
|
||
|
||
|
||
_LOCK_TIMEOUT_SEC = 5.0
|
||
|
||
|
||
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)
|
||
# Ожидание ограничено по времени. Неблокирующий захват был неверен:
|
||
# при неудаче исключение проглатывалось, и запись шла БЕЗ блокировки.
|
||
# Но бесконечное ожидание не лучше: на Unix flock(LOCK_EX) висит
|
||
# вечно, и один застрявший держатель подвесил бы приложение целиком.
|
||
# Секция короткая — запись файла состояния, — поэтому пары секунд
|
||
# с запасом хватает, а дальше честный отказ.
|
||
deadline = time.time() + _LOCK_TIMEOUT_SEC
|
||
while True:
|
||
try:
|
||
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)
|
||
break
|
||
except OSError:
|
||
if time.time() >= deadline:
|
||
raise
|
||
time.sleep(0.05)
|
||
except Exception:
|
||
if self._fd is not None:
|
||
try:
|
||
os.close(self._fd)
|
||
except Exception:
|
||
pass
|
||
self._fd = None
|
||
raise
|
||
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()
|
||
|
||
try:
|
||
from antigravity_provider.router.event_bus import (
|
||
EVENT_ACCOUNT_ADDED,
|
||
EVENT_ACCOUNT_AUTH_CHANGED,
|
||
EventBus,
|
||
)
|
||
EventBus.get().subscribe(EVENT_ACCOUNT_ADDED, self._on_account_event)
|
||
EventBus.get().subscribe(EVENT_ACCOUNT_AUTH_CHANGED, self._on_account_event)
|
||
except Exception:
|
||
pass
|
||
|
||
def _on_account_event(self, _event_name: str, payload: Any) -> None:
|
||
"""Handle account lifecycle events by recovering profile from AUTH_REQUIRED state."""
|
||
profile_id = None
|
||
if isinstance(payload, dict):
|
||
profile_id = payload.get("profile_id")
|
||
elif hasattr(payload, "profile_id"):
|
||
profile_id = getattr(payload, "profile_id")
|
||
if not profile_id:
|
||
return
|
||
profile_id = str(profile_id)
|
||
with self._lock:
|
||
if profile_id in self._profiles:
|
||
rec = self._profiles[profile_id]
|
||
rec.overall_state = HEALTHY
|
||
rec.last_error = None
|
||
rec.auth_error_at = None
|
||
for frec in rec.families.values():
|
||
if frec.state == AUTH_REQUIRED:
|
||
frec.state = HEALTHY
|
||
frec.reset_at = None
|
||
frec.reason = None
|
||
frec.last_error = None
|
||
self._save_state()
|
||
|
||
def _find_profile_auth_files(self, profile_id: str) -> list[Path]:
|
||
"""Locate credentials and auth files for the specified profile."""
|
||
files: list[Path] = []
|
||
try:
|
||
pdir = paths.get_profile_dir(profile_id)
|
||
for cand in (pdir / "auth.json", pdir / ".gemini" / "oauth_creds.json", pdir / "oauth_creds.json"):
|
||
if cand.is_file():
|
||
files.append(cand)
|
||
except Exception:
|
||
pass
|
||
|
||
try:
|
||
hermes_home = paths.get_hermes_home()
|
||
if hermes_home.is_dir():
|
||
for pdir in hermes_home.glob(f"*_profiles/{profile_id}"):
|
||
if pdir.is_dir():
|
||
for cand in (pdir / "auth.json", pdir / ".gemini" / "oauth_creds.json", pdir / "oauth_creds.json"):
|
||
if cand.is_file() and cand not in files:
|
||
files.append(cand)
|
||
except Exception:
|
||
pass
|
||
return files
|
||
|
||
def _check_and_recover_auth(self, record: ProfileHealthRecord) -> bool:
|
||
"""Check if auth files exist, are valid, and were modified after an auth failure."""
|
||
auth_files = self._find_profile_auth_files(record.profile_id)
|
||
if not auth_files:
|
||
return False
|
||
|
||
recovered = False
|
||
for f in auth_files:
|
||
try:
|
||
stat = f.stat()
|
||
if stat.st_size == 0:
|
||
continue
|
||
if record.auth_error_at is not None and stat.st_mtime <= record.auth_error_at:
|
||
continue
|
||
content = json.loads(f.read_text(encoding="utf-8"))
|
||
if isinstance(content, dict) and len(content) > 0:
|
||
recovered = True
|
||
break
|
||
except Exception:
|
||
continue
|
||
|
||
if recovered:
|
||
record.overall_state = HEALTHY
|
||
record.last_error = None
|
||
record.auth_error_at = None
|
||
for frec in record.families.values():
|
||
if frec.state == AUTH_REQUIRED:
|
||
frec.state = HEALTHY
|
||
frec.reset_at = None
|
||
frec.reason = None
|
||
frec.last_error = None
|
||
self._save_state()
|
||
return True
|
||
return False
|
||
|
||
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"),
|
||
auth_error_at=pdata.get("auth_error_at"),
|
||
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,
|
||
"auth_error_at": precord.auth_error_at,
|
||
"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
|
||
|
||
if record.overall_state == AUTH_REQUIRED:
|
||
self._check_and_recover_auth(record)
|
||
|
||
# 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.auth_error_at = None
|
||
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 reconcile_measured_quota(self, profile_id: str, remaining_by_family: Dict[str, float]) -> bool:
|
||
"""Evaluate measured quota against configured threshold and recover when capacity restores.
|
||
|
||
A successful quota read is authoritative for quota exhaustion, but it
|
||
must not erase unrelated authentication or runtime failures.
|
||
"""
|
||
measured = {family: float(value) for family, value in remaining_by_family.items() if value is not None}
|
||
if not measured:
|
||
return False
|
||
|
||
from antigravity_provider.router.settings_service import get_hub_settings
|
||
settings = get_hub_settings()
|
||
threshold = float(settings.get("quota_threshold_percent", 10.0))
|
||
action = str(settings.get("quota_threshold_action", "notify")).strip().lower()
|
||
|
||
with self._lock:
|
||
record = self.get_or_create(profile_id)
|
||
changed = False
|
||
now = time.time()
|
||
|
||
for family, remaining in measured.items():
|
||
family_record = record.families.get(family)
|
||
if remaining <= threshold:
|
||
if action == "switch":
|
||
if family not in record.families:
|
||
record.families[family] = FamilyHealthRecord(family=family)
|
||
frec = record.families[family]
|
||
if frec.state != QUOTA_EXHAUSTED:
|
||
frec.state = QUOTA_EXHAUSTED
|
||
frec.reset_at = now + 1800
|
||
frec.reason = f"Остаток квоты {remaining:.1f}% <= порога {threshold:.1f}%"
|
||
frec.last_error = frec.reason
|
||
changed = True
|
||
elif action == "notify":
|
||
try:
|
||
from antigravity_provider.router.unified_health import EventLogService
|
||
EventLogService.get().log(
|
||
"quota",
|
||
f"Внимание: остаток квоты профиля {profile_id} ({family}) составляет {remaining:.1f}% (порог: {threshold:.1f}%).",
|
||
level="warning",
|
||
)
|
||
except Exception:
|
||
pass
|
||
elif remaining > threshold:
|
||
if family_record and family_record.state == QUOTA_EXHAUSTED:
|
||
family_record.state = HEALTHY
|
||
family_record.reset_at = None
|
||
family_record.reason = None
|
||
family_record.last_error = None
|
||
family_record.simulated = False
|
||
changed = True
|
||
|
||
if all(value > threshold for value in measured.values()) and record.overall_state == QUOTA_EXHAUSTED:
|
||
record.overall_state = HEALTHY
|
||
record.last_error = None
|
||
record.simulated = False
|
||
changed = True
|
||
elif any(value <= threshold for value in measured.values()) and action == "switch":
|
||
if record.overall_state == HEALTHY:
|
||
record.overall_state = QUOTA_EXHAUSTED
|
||
record.last_error = f"Остаток квоты ниже порога {threshold:.1f}%"
|
||
changed = True
|
||
|
||
if changed:
|
||
self._save_state()
|
||
return changed
|
||
|
||
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
|
||
record.auth_error_at = time.time()
|
||
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.last_error = None
|
||
rec.auth_error_at = None
|
||
rec.simulated = False
|
||
for frec in rec.families.values():
|
||
frec.state = HEALTHY
|
||
frec.reset_at = None
|
||
frec.reason = None
|
||
frec.last_error = 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.last_error = None
|
||
record.auth_error_at = None
|
||
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].reason = None
|
||
record.families[family].last_error = None
|
||
record.families[family].simulated = False
|
||
else:
|
||
for frec in record.families.values():
|
||
frec.state = HEALTHY
|
||
frec.reset_at = None
|
||
frec.reason = None
|
||
frec.last_error = None
|
||
frec.simulated = False
|
||
self._save_state()
|