151 lines
4.4 KiB
Python
151 lines
4.4 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import time
|
|
from typing import TYPE_CHECKING
|
|
from urllib.parse import urlparse
|
|
|
|
from yuxi.channels.models import HealthStatus
|
|
from yuxi.utils.logging_config import logger
|
|
|
|
if TYPE_CHECKING:
|
|
from nio import AsyncClient
|
|
|
|
_DEFAULT_PROBE_TIMEOUT = 15
|
|
|
|
_DEVICE_PREFIX_OPENCLAW = "yuxi-bot"
|
|
|
|
_BACKUP_STATES = {
|
|
0: "not_backed_up",
|
|
1: "keys_backed_up_and_verified",
|
|
2: "keys_backed_up_not_verified",
|
|
3: "backup_in_progress",
|
|
4: "backup_error",
|
|
5: "backup_disabled",
|
|
6: "backup_version_mismatch",
|
|
7: "no_backup_data",
|
|
}
|
|
|
|
|
|
async def probe_matrix(client: AsyncClient, timeout_s: float | None = None) -> HealthStatus:
|
|
start = time.monotonic()
|
|
timeout = timeout_s if timeout_s is not None else _DEFAULT_PROBE_TIMEOUT
|
|
|
|
try:
|
|
resp = await asyncio.wait_for(client.versions(), timeout=timeout)
|
|
latency_ms = (time.monotonic() - start) * 1000
|
|
|
|
parsed = urlparse(client.homeserver)
|
|
server_name = parsed.hostname or client.homeserver
|
|
|
|
whoami_info = await _probe_whoami(client, timeout)
|
|
|
|
metadata = {
|
|
"homeserver": client.homeserver,
|
|
"server_name": server_name,
|
|
"user_id": client.user_id,
|
|
"versions": resp.versions if hasattr(resp, "versions") else [],
|
|
"whoami": whoami_info,
|
|
}
|
|
|
|
device_info = _check_device_health(whoami_info)
|
|
metadata["device_health"] = device_info
|
|
|
|
backup_info = await _check_backup_health(client, timeout)
|
|
metadata["backup_health"] = backup_info
|
|
|
|
return HealthStatus(
|
|
status="healthy",
|
|
latency_ms=latency_ms,
|
|
metadata=metadata,
|
|
)
|
|
except TimeoutError:
|
|
latency_ms = (time.monotonic() - start) * 1000
|
|
return HealthStatus(
|
|
status="unhealthy",
|
|
latency_ms=latency_ms,
|
|
last_error=f"Probe timed out after {timeout}s",
|
|
)
|
|
except ConnectionError as e:
|
|
return HealthStatus(
|
|
status="unhealthy",
|
|
last_error=f"Connection refused: {e}",
|
|
)
|
|
except TimeoutError as e:
|
|
return HealthStatus(
|
|
status="unhealthy",
|
|
last_error=f"Connection timeout: {e}",
|
|
)
|
|
except OSError as e:
|
|
return HealthStatus(
|
|
status="unhealthy",
|
|
last_error=f"Network error: {e}",
|
|
)
|
|
except Exception as e:
|
|
return HealthStatus(
|
|
status="degraded",
|
|
last_error=str(e),
|
|
)
|
|
|
|
|
|
async def _probe_whoami(client: AsyncClient, timeout: float) -> dict:
|
|
try:
|
|
resp = await asyncio.wait_for(client.whoami(), timeout=timeout)
|
|
return {
|
|
"user_id": resp.user_id if hasattr(resp, "user_id") else None,
|
|
"device_id": resp.device_id if hasattr(resp, "device_id") else None,
|
|
"is_guest": resp.is_guest if hasattr(resp, "is_guest") else False,
|
|
}
|
|
except Exception as e:
|
|
logger.debug(f"Matrix whoami probe failed: {e}")
|
|
return {"error": str(e)}
|
|
|
|
|
|
def _check_device_health(whoami: dict) -> dict[str, str]:
|
|
device_id = whoami.get("device_id", "")
|
|
if not device_id:
|
|
return {"status": "unknown", "reason": "no_device_id"}
|
|
|
|
if device_id.startswith(_DEVICE_PREFIX_OPENCLAW):
|
|
return {
|
|
"status": "healthy",
|
|
"device_id": device_id,
|
|
"is_known_prefix": "true",
|
|
}
|
|
|
|
return {
|
|
"status": "unrecognized",
|
|
"device_id": device_id,
|
|
"is_known_prefix": "false",
|
|
}
|
|
|
|
|
|
async def _check_backup_health(client: AsyncClient, timeout: float) -> dict[str, str]:
|
|
try:
|
|
await asyncio.wait_for(client.keys_upload(), timeout=timeout)
|
|
return {"status": "ok", "message": "key_upload_reachable"}
|
|
except Exception:
|
|
pass
|
|
|
|
try:
|
|
await asyncio.wait_for(client.keys_query(), timeout=timeout)
|
|
return {"status": "ok", "message": "keys_query_reachable"}
|
|
except Exception:
|
|
pass
|
|
|
|
try:
|
|
await asyncio.wait_for(client.query_keys(), timeout=timeout)
|
|
return {"status": "ok", "message": "query_keys_reachable"}
|
|
except Exception:
|
|
pass
|
|
|
|
return {"status": "unknown", "message": "backup_status_undetermined"}
|
|
|
|
|
|
def describe_backup_state(code: int) -> str:
|
|
return _BACKUP_STATES.get(code, f"unknown_state_{code}")
|
|
|
|
|
|
def is_backup_healthy(state_code: int) -> bool:
|
|
return state_code in (1, 2)
|