463 lines
17 KiB
Python
463 lines
17 KiB
Python
"""Self-healing availability guard for persistent Online Agent workers.
|
|
|
|
The guard never claims Tasks and never writes Task/Job/Git databases. It only:
|
|
- ensures required worker/executor systemd units are enabled + active;
|
|
- checks worker-owned liveness JSON written by the real poll/heartbeat loop;
|
|
- restarts an active worker only when its liveness is stale;
|
|
- honors an explicit maintenance inhibit file;
|
|
- writes a bounded status JSON under /run for observability.
|
|
|
|
Idle liveness is the successful poll loop itself. Busy liveness is refreshed by the
|
|
Task lease-heartbeat thread, so long-running Tasks are not mistaken for hung workers.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
from dataclasses import asdict, dataclass
|
|
import argparse
|
|
import json
|
|
import os
|
|
from pathlib import Path
|
|
import shlex
|
|
import subprocess
|
|
import tempfile
|
|
import time
|
|
from typing import Any, Callable, Mapping, Sequence
|
|
|
|
SCHEMA = "aios.worker-availability.status.v1"
|
|
DEFAULT_WORKER_UNITS = (
|
|
"aios-online-agent-worker-20260815.service",
|
|
"aios-hermes-online-agent-worker.service",
|
|
"aios-openclaw-online-agent-worker.service",
|
|
)
|
|
DEFAULT_DEPENDENCY_UNITS = (
|
|
"aios-external-agent-executor.service",
|
|
)
|
|
DEFAULT_RUN_ROOT = Path("/run/aios-worker-availability")
|
|
DEFAULT_STATUS = DEFAULT_RUN_ROOT / "status.json"
|
|
DEFAULT_RECOVERY = DEFAULT_RUN_ROOT / "recovery.json"
|
|
DEFAULT_INHIBIT = DEFAULT_RUN_ROOT / "maintenance.lock"
|
|
DEFAULT_MAX_STALE_SECONDS = 90
|
|
DEFAULT_RESTART_COOLDOWN_SECONDS = 60
|
|
DEFAULT_STARTUP_WAIT_SECONDS = 10
|
|
ACCEPTED_ENABLED_STATES = frozenset({"enabled", "enabled-runtime", "static", "indirect", "generated"})
|
|
|
|
|
|
class AvailabilityGuardError(RuntimeError):
|
|
pass
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class UnitSnapshot:
|
|
unit: str
|
|
role: str
|
|
load_state: str
|
|
active_state: str
|
|
unit_file_state: str
|
|
enabled_ok: bool
|
|
liveness_path: str | None = None
|
|
liveness_state: str | None = None
|
|
last_seen_epoch: float | None = None
|
|
liveness_age_seconds: float | None = None
|
|
fresh: bool | None = None
|
|
action: str = "NONE"
|
|
detail: str | None = None
|
|
|
|
|
|
class Systemd:
|
|
def __init__(
|
|
self,
|
|
runner: Callable[..., subprocess.CompletedProcess[str]] | None = None,
|
|
) -> None:
|
|
self.runner = runner or subprocess.run
|
|
|
|
def _run(self, args: Sequence[str], *, check: bool = False) -> subprocess.CompletedProcess[str]:
|
|
result = self.runner(
|
|
["/usr/bin/systemctl", *args],
|
|
text=True,
|
|
capture_output=True,
|
|
check=False,
|
|
timeout=20,
|
|
)
|
|
if check and result.returncode != 0:
|
|
raise AvailabilityGuardError(
|
|
"SYSTEMD_COMMAND_FAILED:"
|
|
+ ":".join(args)
|
|
+ ":"
|
|
+ str(result.returncode)
|
|
)
|
|
return result
|
|
|
|
def property(self, unit: str, name: str) -> str:
|
|
result = self._run(["show", unit, f"--property={name}", "--value"])
|
|
if result.returncode != 0:
|
|
return ""
|
|
return result.stdout.strip()
|
|
|
|
def environment(self, unit: str) -> dict[str, str]:
|
|
text = self.property(unit, "Environment")
|
|
result: dict[str, str] = {}
|
|
for item in shlex.split(text):
|
|
if "=" not in item:
|
|
continue
|
|
key, value = item.split("=", 1)
|
|
result[key] = value
|
|
return result
|
|
|
|
def enable(self, unit: str) -> None:
|
|
self._run(["enable", unit], check=True)
|
|
|
|
def start(self, unit: str) -> None:
|
|
self._run(["start", unit], check=True)
|
|
|
|
def restart(self, unit: str) -> None:
|
|
self._run(["restart", unit], check=True)
|
|
|
|
|
|
def _atomic_json(path: Path, value: Mapping[str, Any]) -> None:
|
|
path.parent.mkdir(parents=True, exist_ok=True, mode=0o755)
|
|
if path.parent.is_symlink() or path.is_symlink():
|
|
raise AvailabilityGuardError("STATUS_PATH_SYMLINK_REJECTED")
|
|
fd, temporary = tempfile.mkstemp(prefix=f".{path.name}.", dir=path.parent)
|
|
try:
|
|
os.fchmod(fd, 0o644)
|
|
with os.fdopen(fd, "w", encoding="utf-8") as handle:
|
|
json.dump(value, handle, ensure_ascii=False, sort_keys=True, indent=2)
|
|
handle.write("\n")
|
|
handle.flush()
|
|
os.fsync(handle.fileno())
|
|
os.replace(temporary, path)
|
|
os.chmod(path, 0o644)
|
|
finally:
|
|
if os.path.exists(temporary):
|
|
os.unlink(temporary)
|
|
|
|
|
|
def _load_json(path: Path) -> dict[str, Any] | None:
|
|
if path.is_symlink() or not path.is_file():
|
|
return None
|
|
try:
|
|
value = json.loads(path.read_text(encoding="utf-8"))
|
|
except (OSError, json.JSONDecodeError):
|
|
return None
|
|
return dict(value) if isinstance(value, Mapping) else None
|
|
|
|
|
|
def _state_root(environment: Mapping[str, str], unit: str) -> Path:
|
|
explicit = environment.get("AIOS_WORKER_STATE_ROOT")
|
|
if explicit:
|
|
return Path(explicit)
|
|
fallbacks = {
|
|
"aios-online-agent-worker-20260815.service": Path(
|
|
"/var/lib/aios-online-agent-worker-20260815"
|
|
),
|
|
"aios-hermes-online-agent-worker.service": Path(
|
|
"/var/lib/aios-online-agent-hermes"
|
|
),
|
|
"aios-openclaw-online-agent-worker.service": Path(
|
|
"/var/lib/aios-online-agent-openclaw"
|
|
),
|
|
}
|
|
if unit not in fallbacks:
|
|
raise AvailabilityGuardError("WORKER_STATE_ROOT_UNAVAILABLE:" + unit)
|
|
return fallbacks[unit]
|
|
|
|
|
|
def _liveness_path(environment: Mapping[str, str], unit: str) -> Path:
|
|
explicit = environment.get("AIOS_WORKER_LIVENESS_PATH")
|
|
if explicit:
|
|
return Path(explicit)
|
|
return _state_root(environment, unit) / "worker-liveness.json"
|
|
|
|
|
|
class AvailabilityGuard:
|
|
def __init__(
|
|
self,
|
|
*,
|
|
systemd: Systemd | None = None,
|
|
worker_units: Sequence[str] = DEFAULT_WORKER_UNITS,
|
|
dependency_units: Sequence[str] = DEFAULT_DEPENDENCY_UNITS,
|
|
status_path: Path = DEFAULT_STATUS,
|
|
recovery_path: Path = DEFAULT_RECOVERY,
|
|
inhibit_path: Path = DEFAULT_INHIBIT,
|
|
max_stale_seconds: int = DEFAULT_MAX_STALE_SECONDS,
|
|
restart_cooldown_seconds: int = DEFAULT_RESTART_COOLDOWN_SECONDS,
|
|
startup_wait_seconds: int = DEFAULT_STARTUP_WAIT_SECONDS,
|
|
now_fn: Callable[[], float] = time.time,
|
|
sleep_fn: Callable[[float], None] = time.sleep,
|
|
) -> None:
|
|
self.systemd = systemd or Systemd()
|
|
self.worker_units = tuple(worker_units)
|
|
self.dependency_units = tuple(dependency_units)
|
|
self.status_path = Path(status_path)
|
|
self.recovery_path = Path(recovery_path)
|
|
self.inhibit_path = Path(inhibit_path)
|
|
self.max_stale_seconds = int(max_stale_seconds)
|
|
self.restart_cooldown_seconds = int(restart_cooldown_seconds)
|
|
self.startup_wait_seconds = int(startup_wait_seconds)
|
|
self.now_fn = now_fn
|
|
self.sleep_fn = sleep_fn
|
|
if self.max_stale_seconds < 15:
|
|
raise AvailabilityGuardError("MAX_STALE_SECONDS_TOO_SMALL")
|
|
if self.restart_cooldown_seconds < 15:
|
|
raise AvailabilityGuardError("RESTART_COOLDOWN_TOO_SMALL")
|
|
if self.startup_wait_seconds < 0 or self.startup_wait_seconds > 60:
|
|
raise AvailabilityGuardError("STARTUP_WAIT_INVALID")
|
|
|
|
def _recovery_state(self) -> dict[str, float]:
|
|
value = _load_json(self.recovery_path) or {}
|
|
raw = value.get("last_restart_epoch")
|
|
if not isinstance(raw, Mapping):
|
|
return {}
|
|
result: dict[str, float] = {}
|
|
for unit, epoch in raw.items():
|
|
try:
|
|
result[str(unit)] = float(epoch)
|
|
except (TypeError, ValueError):
|
|
continue
|
|
return result
|
|
|
|
def _save_recovery_state(self, state: Mapping[str, float]) -> None:
|
|
_atomic_json(
|
|
self.recovery_path,
|
|
{
|
|
"schema": "aios.worker-availability.recovery.v1",
|
|
"last_restart_epoch": dict(state),
|
|
},
|
|
)
|
|
|
|
def _basic_snapshot(self, unit: str, role: str) -> tuple[str, str, str, bool]:
|
|
load = self.systemd.property(unit, "LoadState")
|
|
active = self.systemd.property(unit, "ActiveState")
|
|
enabled = self.systemd.property(unit, "UnitFileState")
|
|
return load, active, enabled, enabled in ACCEPTED_ENABLED_STATES
|
|
|
|
def _ensure_enabled_active(self, unit: str, role: str) -> tuple[str, bool]:
|
|
load, active, enabled, enabled_ok = self._basic_snapshot(unit, role)
|
|
if load != "loaded":
|
|
raise AvailabilityGuardError("REQUIRED_UNIT_NOT_LOADED:" + unit)
|
|
action = "NONE"
|
|
if not enabled_ok:
|
|
self.systemd.enable(unit)
|
|
action = "ENABLE"
|
|
if active != "active":
|
|
self.systemd.start(unit)
|
|
action = "START" if action == "NONE" else "ENABLE_START"
|
|
return action, active == "active"
|
|
|
|
def _wait_for_liveness(self, path: Path) -> dict[str, Any] | None:
|
|
deadline = self.now_fn() + self.startup_wait_seconds
|
|
while True:
|
|
value = _load_json(path)
|
|
if value is not None:
|
|
try:
|
|
epoch = float(value.get("last_seen_epoch"))
|
|
except (TypeError, ValueError):
|
|
epoch = 0.0
|
|
if epoch > 0 and self.now_fn() - epoch <= self.max_stale_seconds:
|
|
return value
|
|
if self.now_fn() >= deadline:
|
|
return value
|
|
self.sleep_fn(0.5)
|
|
|
|
def run(self, *, heal: bool) -> dict[str, Any]:
|
|
now = self.now_fn()
|
|
if self.inhibit_path.exists():
|
|
result = {
|
|
"schema": SCHEMA,
|
|
"status": "MAINTENANCE_INHIBITED",
|
|
"generated_at_epoch": now,
|
|
"maintenance_inhibit": str(self.inhibit_path),
|
|
"workers": [],
|
|
"dependencies": [],
|
|
"all_fresh": False,
|
|
"actions": [],
|
|
}
|
|
_atomic_json(self.status_path, result)
|
|
return result
|
|
|
|
actions: list[dict[str, str]] = []
|
|
dependencies: list[UnitSnapshot] = []
|
|
for unit in self.dependency_units:
|
|
load, active, enabled, enabled_ok = self._basic_snapshot(unit, "dependency")
|
|
action = "NONE"
|
|
if heal:
|
|
if load != "loaded":
|
|
raise AvailabilityGuardError("REQUIRED_UNIT_NOT_LOADED:" + unit)
|
|
if not enabled_ok:
|
|
self.systemd.enable(unit)
|
|
action = "ENABLE"
|
|
if active != "active":
|
|
self.systemd.start(unit)
|
|
action = "START" if action == "NONE" else "ENABLE_START"
|
|
if action != "NONE":
|
|
actions.append({"unit": unit, "action": action})
|
|
load, active, enabled, enabled_ok = self._basic_snapshot(unit, "dependency")
|
|
dependencies.append(
|
|
UnitSnapshot(
|
|
unit=unit,
|
|
role="dependency",
|
|
load_state=load,
|
|
active_state=active,
|
|
unit_file_state=enabled,
|
|
enabled_ok=enabled_ok,
|
|
fresh=(active == "active"),
|
|
action=action,
|
|
)
|
|
)
|
|
|
|
recovery = self._recovery_state()
|
|
workers: list[UnitSnapshot] = []
|
|
for unit in self.worker_units:
|
|
load, active, enabled, enabled_ok = self._basic_snapshot(unit, "worker")
|
|
environment = self.systemd.environment(unit) if load == "loaded" else {}
|
|
path = _liveness_path(environment, unit) if load == "loaded" else None
|
|
action = "NONE"
|
|
|
|
if heal:
|
|
if load != "loaded":
|
|
raise AvailabilityGuardError("REQUIRED_UNIT_NOT_LOADED:" + unit)
|
|
if not enabled_ok:
|
|
self.systemd.enable(unit)
|
|
action = "ENABLE"
|
|
if active != "active":
|
|
self.systemd.start(unit)
|
|
action = "START" if action == "NONE" else "ENABLE_START"
|
|
actions.append({"unit": unit, "action": action})
|
|
if path is not None:
|
|
self._wait_for_liveness(path)
|
|
load, active, enabled, enabled_ok = self._basic_snapshot(unit, "worker")
|
|
|
|
live = _load_json(path) if path is not None else None
|
|
last_seen: float | None = None
|
|
state: str | None = None
|
|
age: float | None = None
|
|
fresh = False
|
|
detail: str | None = None
|
|
if live is not None:
|
|
state = str(live.get("state") or "")
|
|
try:
|
|
last_seen = float(live.get("last_seen_epoch"))
|
|
except (TypeError, ValueError):
|
|
last_seen = None
|
|
if last_seen is not None:
|
|
age = max(0.0, now - last_seen)
|
|
fresh = age <= self.max_stale_seconds
|
|
if live.get("agent_id") in (None, ""):
|
|
fresh = False
|
|
detail = "LIVENESS_AGENT_ID_MISSING"
|
|
else:
|
|
detail = "LIVENESS_MISSING"
|
|
|
|
if heal and active == "active" and not fresh:
|
|
previous = float(recovery.get(unit, 0.0))
|
|
if now - previous >= self.restart_cooldown_seconds:
|
|
self.systemd.restart(unit)
|
|
recovery[unit] = now
|
|
action = "RESTART_STALE"
|
|
actions.append({"unit": unit, "action": action})
|
|
if path is not None:
|
|
live = self._wait_for_liveness(path)
|
|
load, active, enabled, enabled_ok = self._basic_snapshot(unit, "worker")
|
|
if live is not None:
|
|
state = str(live.get("state") or "")
|
|
try:
|
|
last_seen = float(live.get("last_seen_epoch"))
|
|
except (TypeError, ValueError):
|
|
last_seen = None
|
|
if last_seen is not None:
|
|
age = max(0.0, self.now_fn() - last_seen)
|
|
fresh = age <= self.max_stale_seconds
|
|
detail = None if fresh else "LIVENESS_STALE_AFTER_RESTART"
|
|
else:
|
|
detail = "RESTART_COOLDOWN_ACTIVE"
|
|
|
|
workers.append(
|
|
UnitSnapshot(
|
|
unit=unit,
|
|
role="worker",
|
|
load_state=load,
|
|
active_state=active,
|
|
unit_file_state=enabled,
|
|
enabled_ok=enabled_ok,
|
|
liveness_path=str(path) if path is not None else None,
|
|
liveness_state=state,
|
|
last_seen_epoch=last_seen,
|
|
liveness_age_seconds=age,
|
|
fresh=bool(active == "active" and enabled_ok and fresh),
|
|
action=action,
|
|
detail=detail,
|
|
)
|
|
)
|
|
|
|
if heal:
|
|
self._save_recovery_state(recovery)
|
|
|
|
all_fresh = (
|
|
all(item.active_state == "active" and item.enabled_ok and item.fresh for item in workers)
|
|
and all(item.active_state == "active" for item in dependencies)
|
|
)
|
|
result = {
|
|
"schema": SCHEMA,
|
|
"status": "HEALTHY" if all_fresh else "DEGRADED",
|
|
"generated_at_epoch": self.now_fn(),
|
|
"max_stale_seconds": self.max_stale_seconds,
|
|
"maintenance_inhibit": str(self.inhibit_path),
|
|
"workers": [asdict(item) for item in workers],
|
|
"dependencies": [asdict(item) for item in dependencies],
|
|
"all_fresh": all_fresh,
|
|
"actions": actions,
|
|
}
|
|
_atomic_json(self.status_path, result)
|
|
return result
|
|
|
|
|
|
def main() -> int:
|
|
parser = argparse.ArgumentParser()
|
|
mode = parser.add_mutually_exclusive_group(required=True)
|
|
mode.add_argument("--heal", action="store_true")
|
|
mode.add_argument("--check", action="store_true")
|
|
parser.add_argument("--status-file", default=str(DEFAULT_STATUS))
|
|
parser.add_argument("--recovery-file", default=str(DEFAULT_RECOVERY))
|
|
parser.add_argument("--inhibit-file", default=str(DEFAULT_INHIBIT))
|
|
parser.add_argument(
|
|
"--max-stale-seconds",
|
|
type=int,
|
|
default=int(
|
|
os.environ.get(
|
|
"AIOS_WORKER_AVAILABILITY_MAX_STALE_SECONDS",
|
|
str(DEFAULT_MAX_STALE_SECONDS),
|
|
)
|
|
),
|
|
)
|
|
parser.add_argument(
|
|
"--restart-cooldown-seconds",
|
|
type=int,
|
|
default=int(
|
|
os.environ.get(
|
|
"AIOS_WORKER_AVAILABILITY_RESTART_COOLDOWN_SECONDS",
|
|
str(DEFAULT_RESTART_COOLDOWN_SECONDS),
|
|
)
|
|
),
|
|
)
|
|
args = parser.parse_args()
|
|
guard = AvailabilityGuard(
|
|
status_path=Path(args.status_file),
|
|
recovery_path=Path(args.recovery_file),
|
|
inhibit_path=Path(args.inhibit_file),
|
|
max_stale_seconds=args.max_stale_seconds,
|
|
restart_cooldown_seconds=args.restart_cooldown_seconds,
|
|
)
|
|
try:
|
|
result = guard.run(heal=args.heal)
|
|
except AvailabilityGuardError as exc:
|
|
print(json.dumps({"status": "FAILED", "error": str(exc)}, sort_keys=True))
|
|
return 20
|
|
print(json.dumps(result, ensure_ascii=False, sort_keys=True))
|
|
if args.check and not result.get("all_fresh"):
|
|
return 10
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|