herdr-session-switch/herdr_hub/restart.py

"""Standalone POSIX update/handoff helper, also sent to remote Python over SSH."""
from __future__ import annotations

import json
import os
from pathlib import Path
import socket
import subprocess
import sys
import time
import uuid


def request(socket_path: str, method: str, params: dict, timeout: float = 10) -> dict:
    request_id = uuid.uuid4().hex
    with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as connection:
        connection.settimeout(timeout)
        connection.connect(socket_path)
        connection.sendall((json.dumps({"id": request_id, "method": method, "params": params}) + "\n").encode())
        with connection.makefile("rb") as stream:
            raw = stream.readline(1024 * 1024)
    if not raw.endswith(b"\n"):
        raise RuntimeError("Herdr 控制响应不完整；请检查 Session 状态后再重试")
    response = json.loads(raw)
    if not isinstance(response, dict) or response.get("id") != request_id:
        raise RuntimeError("Herdr 控制响应 ID 不匹配")
    if response.get("error"):
        raise RuntimeError(json.dumps(response["error"], ensure_ascii=False))
    result = response.get("result")
    if not isinstance(result, dict):
        raise RuntimeError("Herdr 控制响应缺少结果")
    return result


def ping_when_reachable(socket_path: str, *, timeout: float = 20) -> dict:
    """Ping a server that may still be restarting right after an update.

    The native updater can return while the new server is still binding the
    socket; a single refused connection must not be reported as a failed
    update. The first attempt is immediate so a healthy server adds no delay.
    """
    deadline: float | None = None
    while True:
        try:
            return request(socket_path, "ping", {}, timeout=2)
        except (OSError, ValueError, RuntimeError) as error:
            now = time.monotonic()
            if deadline is None:
                deadline = now + timeout
            elif now >= deadline:
                raise RuntimeError(
                    f"更新后 {timeout:g} 秒内无法连接 Herdr 服务器：{error}；"
                    "请检查 Session 状态，不会自动重启"
                ) from error
            time.sleep(.2)


def installed_version(binary: str) -> str:
    version = subprocess.check_output([binary, "--version"], text=True, timeout=10).strip()
    if not version.startswith("herdr ") or not version[6:].strip():
        raise RuntimeError("无法识别已安装的 Herdr 版本：" + version)
    return version[6:].strip()


def binary_identity(binary: str) -> tuple[int, int, int, int]:
    info = os.stat(binary)
    return info.st_dev, info.st_ino, info.st_size, info.st_mtime_ns


def update_and_handoff(binary: str, session: str) -> None:
    binary = str(Path(binary).absolute())
    status = subprocess.check_output(
        [binary, "--session", session, "status", "server"], text=True, timeout=10,
    )
    fields = dict(line.split(": ", 1) for line in status.splitlines() if ": " in line)
    socket_path = fields.get("socket", "")
    if fields.get("status") != "running" or not os.path.isabs(socket_path):
        raise RuntimeError("当前 Session 没有可用的运行中服务器，不能执行热切换")
    before = request(socket_path, "ping", {})
    if before.get("capabilities", {}).get("live_handoff") is not True:
        raise RuntimeError("当前服务器不支持热切换；可先使用“更新 Herdr”，再安排普通重启")
    print(f"更新目标 Session：{session}；当前服务器：{before.get('version', '未知')}", flush=True)
    identity_before = binary_identity(binary)
    # 0.8.x plain update refuses to complete with live sessions and no TTY.
    # Let the native updater perform its opt-in handoff, scoped to this session.
    subprocess.run([binary, "--session", session, "update", "--handoff"], check=True, timeout=300)
    version = installed_version(binary)
    current = ping_when_reachable(socket_path)
    if current.get("version") != version:
        if binary_identity(binary) != identity_before:
            raise RuntimeError(
                f"已安装 {version}，但原生热切换后服务器仍是 {current.get('version', '未知')}；"
                "不会重复发送热切换请求或强制重启，请检查上方原生命令输出"
            )
        print(f"正在热切换服务器到 {version}，保留运行中的进程…", flush=True)
        # An unchanged binary means the updater may have returned 'up to date'
        # before touching an older running server. Only then send a handoff.
        request(socket_path, "server.live_handoff", {
            "import_exe": str(Path(binary).resolve()), "expected_version": version,
        }, timeout=40)
        deadline = time.monotonic() + 20
        while True:
            try:
                current = request(socket_path, "ping", {}, timeout=2)
                if current.get("version") == version:
                    break
            except (OSError, ValueError, RuntimeError):
                pass
            if time.monotonic() >= deadline:
                raise RuntimeError("热切换后未确认新服务器就绪；请检查状态，不会自动停止 Session")
            time.sleep(.2)
    print(f"服务器已确认使用 {version}，可以重新连接。", flush=True)


if __name__ == "__main__":
    try:
        update_and_handoff(sys.argv[1], sys.argv[2])
    except Exception as error:
        print(f"更新/热切换未完成：{error}\n没有自动执行普通重启；请检查当前 Session。", file=sys.stderr, flush=True)
        raise SystemExit(1)