feat(host_agent): Steward hands a host a maintenance job and follows it through the reboot (#5767)

#5767 "Steward hands a host a maintenance job and follows it to the end,
through its own restarts", milestone 514.

A host_maintenance_jobs row per job: queued → updating → [draining →]
rebooting → [undraining →] done, or failed / timed_out. Steward's answer to
an ingest post carries the step that host's agent should carry out now; the
agent (1.18.0) leaves each step id for the root helper once, and the helper
refuses an id it already ran, so repeats are harmless. Results come back as
the #5766 maintenance events, matched by step id; the reboot completes when
the host's `started` event has a boot_at after the reboot was handed out.

The job reboots only when the update says one is needed (or never, if asked).
A swarm node is drained first by itself if it's a manager or by a manager
with the helper, and made active after; with no manager to drain it the job
is refused. A minute tick times out a job past its state's deadline. A
maintenance window over host_agent/<host> lasts the rest of the job.

Everything is on the row, so a Steward restart, or the reboot of the host
Steward runs on, resumes from the next post. Ingest runs it in its own
savepoint so a job going wrong never costs the samples.

The Agent panel leads with a running job and says what a finished one did;
Manage agent has Update now (with "never reboot"), or says why a job can't
start. Integration test: hand off, Steward restarts, the post-boot events
close the job.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

plugins/host_agent/maintenance.py

@@ -0,0 +1,384 @@
+# plugins/host_agent/maintenance.py
+"""A host's maintenance job, followed to its end through any restart (#5767).
+
+Steward plans and the host's agent acts: a job is one host's update and, only
+when the update says one is needed, its reboot, with the host's swarm node
+drained first and made active again after. Each action is one step the agent
+hands to its root helper (maint.py, #5766). Everything lives on the job's row,
+so Steward restarting, or the host Steward runs on rebooting, loses nothing:
+the next post from the agent carries the step's result, and the `started`
+event a host sends after booting is what completes a reboot.
+
+The transitions here are model-free (a job is anything with the attributes),
+so they're unit-tested without a database; routes.ingest and the scheduler's
+tick apply them.
+
+States: queued → updating → [draining →] rebooting → [undraining →] done,
+or failed / timed_out. A job with no reboot needed goes updating → done.
+"""
+from __future__ import annotations
+
+from datetime import datetime, timedelta, timezone
+
+ACTIVE = ("queued", "updating", "draining", "rebooting", "undraining")
+FINISHED = ("done", "failed", "timed_out")
+# The helper action each handed state asks for.
+STEP_ACTION = {"updating": "update", "draining": "drain", "rebooting": "reboot",
+               "undraining": "undrain"}
+# How long a state may last before the job is given up. Queued counts from the
+# job's creation (the host may be silent); the rest from their hand-off.
+DEADLINES = {
+    "queued": timedelta(hours=1),
+    "updating": timedelta(minutes=90),     # the helper itself stops at 60
+    "draining": timedelta(minutes=10),
+    "rebooting": timedelta(minutes=30),
+    "undraining": timedelta(minutes=10),
+}
+# A host's clock and Steward's differ by seconds; a boot this far before the
+# reboot was handed out is still taken as that reboot.
+BOOT_SLACK = timedelta(seconds=60)
+DETAIL_MAX = 500
+
+
+def step_id(job, state: str) -> str:
+    """The helper request id for a job's step: unique per job and step, and in
+    the helper's id alphabet (letters, digits, dashes)."""
+    return f"{job.id}-{STEP_ACTION[state]}"
+
+
+def _enter(job, state: str, now: datetime) -> None:
+    job.state, job.state_at = state, now
+    job.step_id = step_id(job, state) if state in STEP_ACTION else None
+    job.step_seen = False               # a new step is handed out until its result arrives
+    if state in FINISHED:
+        job.finished_at = now
+
+
+def _finish(job, state: str, now: datetime, detail: str | None = None) -> None:
+    if detail:
+        job.detail = detail[:DETAIL_MAX]
+    _enter(job, state, now)
+
+
+def start(job, now: datetime) -> None:
+    """The host's agent is about to be handed the job: update first."""
+    if job.state == "queued":
+        _enter(job, "updating", now)
+
+
+def executor(job) -> str:
+    """The host whose agent carries out the current step: the job's own host,
+    except a drain or undrain, which a swarm manager does."""
+    if job.state in ("draining", "undraining") and job.drain_host_id:
+        return job.drain_host_id
+    return job.host_id
+
+
+def handoff(job) -> dict | None:
+    """What to hand the executing agent now, or None. Handed on every post
+    until a result for the step arrives; the agent and the helper both ignore
+    an id they've already taken, so a repeat is harmless."""
+    if job.state not in STEP_ACTION or job.step_seen:
+        return None
+    out = {"id": job.step_id, "action": STEP_ACTION[job.state]}
+    if job.state in ("draining", "undraining") and job.node_id:
+        out["node"] = job.node_id
+    return out
+
+
+def _after_update(job, now: datetime) -> None:
+    if not (job.reboot_needed and job.reboot_policy != "never"):
+        _enter(job, "done", now)
+    elif job.node_id and job.drain_host_id:
+        _enter(job, "draining", now)
+    else:
+        _enter(job, "rebooting", now)
+
+
+def apply_result(job, event: dict, now: datetime) -> bool:
+    """Apply a `maintenance` event for this job's current step; True when the
+    job changed. A result for a step the job has moved past is ignored, so a
+    repeated or late report does nothing."""
+    if job.state not in STEP_ACTION or event.get("id") != job.step_id:
+        return False
+    state = event.get("state")
+    if state in ("running", "rebooting"):
+        if job.step_seen:
+            return False
+        job.step_seen = True
+        return True
+    if state not in ("done", "failed"):
+        return False
+    if state == "failed" or event.get("ok") is False:
+        what = STEP_ACTION[job.state]
+        detail = f"{what} failed: {event.get('detail') or 'no detail'}"
+        if job.state == "undraining":
+            detail += " (the node is still drained)"
+        _finish(job, "failed", now, detail)
+        return True
+    if job.state == "updating":
+        if isinstance(event.get("packages"), int):
+            job.packages = event["packages"]
+        job.reboot_needed = event.get("reboot_required") is True
+        _after_update(job, now)
+    elif job.state == "draining":
+        _enter(job, "rebooting", now)
+    elif job.state == "undraining":
+        _enter(job, "done", now)
+    else:                               # rebooting: the boot completes it, not this
+        return False
+    return True
+
+
+def apply_boot(job, boot_at: datetime | None, now: datetime) -> bool:
+    """A `started` event from the job's host. A boot after the reboot was
+    handed out is that reboot: the host is back."""
+    if job.state != "rebooting" or boot_at is None or job.state_at is None:
+        return False
+    if boot_at < job.state_at - BOOT_SLACK:
+        return False                    # the agent restarted; the host didn't
+    job.back_at = now
+    if job.node_id and job.drain_host_id:
+        _enter(job, "undraining", now)
+    else:
+        _enter(job, "done", now)
+    return True
+
+
+def check_deadline(job, now: datetime) -> bool:
+    """Give up on a job whose state has lasted past its deadline."""
+    limit = DEADLINES.get(job.state)
+    since = job.state_at or job.created_at
+    if limit is None or since is None or now - since <= limit:
+        return False
+    waiting = {"queued": "the host's agent didn't check in",
+               "updating": "the update didn't finish",
+               "draining": "the node wasn't drained",
+               "rebooting": "the host didn't come back",
+               "undraining": "the node wasn't made active again"}[job.state]
+    _finish(job, "timed_out", now, f"{waiting} within {int(limit.total_seconds() // 60)} min")
+    return True
+
+
+def window_end(job, now: datetime) -> datetime:
+    """When the job's maintenance window should end: the longest the rest of
+    the job could take, or now once it has finished."""
+    if job.state in FINISHED:
+        return now
+    left = sum((DEADLINES[s] for s in ACTIVE[ACTIVE.index(job.state):]), timedelta())
+    return now + left
+
+
+def utcnow() -> datetime:
+    return datetime.now(timezone.utc)
+
+
+# ── applied to the database: ingest, the tick, and starting a job ────────────
+
+async def _window(session, job, host_name: str, now: datetime) -> None:
+    """Keep the job's maintenance window over the host until the job ends: one
+    scoped to this host's agent metrics, so an alert the job causes waits."""
+    from steward.models.alerts import MaintenanceWindow
+    if job.window_id:
+        w = await session.get(MaintenanceWindow, job.window_id)
+        if w is not None:
+            w.end_at = max(w.start_at, window_end(job, now))
+        return
+    if job.created_by is None or job.state in FINISHED:
+        return                          # a window needs an owner; none for a finished job
+    w = MaintenanceWindow(name=f"Maintenance: {host_name}"[:128], start_at=now,
+                          end_at=window_end(job, now), scope_source="host_agent",
+                          scope_resource=host_name, created_by=job.created_by)
+    session.add(w)
+    await session.flush()
+    job.window_id = w.id
+
+
+async def _host_name(session, host_id: str) -> str:
+    from steward.models.hosts import Host
+    host = await session.get(Host, host_id)
+    return host.name if host is not None else host_id
+
+
+async def apply_events(session, host_id: str, events, now: datetime) -> int:
+    """Apply one post's `maintenance` results and `started` boots to the jobs
+    they belong to; returns how many jobs changed. A result is matched by its
+    step id, so a drain reported by a manager reaches the job of the node it
+    drained."""
+    from sqlalchemy import select
+    from .downtime import _at
+    from .models import HostMaintenanceJob
+    if not isinstance(events, list):
+        return 0
+    changed = set()
+    for event in events:
+        if not isinstance(event, dict):
+            continue
+        if event.get("kind") == "maintenance" and isinstance(event.get("id"), str):
+            job = (await session.execute(select(HostMaintenanceJob).where(
+                HostMaintenanceJob.step_id == event["id"],
+                HostMaintenanceJob.state.in_(ACTIVE)))).scalar_one_or_none()
+            if job is not None and apply_result(job, event, now):
+                changed.add(job)
+        elif event.get("kind") == "started":
+            boot_at = _at(event.get("boot_at"))
+            for job in (await session.execute(select(HostMaintenanceJob).where(
+                    HostMaintenanceJob.host_id == host_id,
+                    HostMaintenanceJob.state == "rebooting"))).scalars():
+                if apply_boot(job, boot_at, now):
+                    changed.add(job)
+    for job in changed:
+        await _window(session, job, await _host_name(session, job.host_id), now)
+    return len(changed)
+
+
+async def handoffs(session, host_id: str, host_name: str, now: datetime) -> list[dict]:
+    """The steps this host's agent should carry out now. A queued job for the
+    host starts here, when its agent is known to be listening."""
+    from sqlalchemy import or_, select
+    from .models import HostMaintenanceJob
+    jobs = (await session.execute(select(HostMaintenanceJob).where(
+        HostMaintenanceJob.state.in_(ACTIVE),
+        or_(HostMaintenanceJob.host_id == host_id,
+            HostMaintenanceJob.drain_host_id == host_id)
+    ).order_by(HostMaintenanceJob.created_at))).scalars().all()
+    out = []
+    for job in jobs:
+        if job.state == "queued" and job.host_id == host_id:
+            start(job, now)
+            await _window(session, job, host_name, now)
+        if executor(job) == host_id:
+            step = handoff(job)
+            if step is not None:
+                out.append(step)
+    return out
+
+
+async def tick(session, now: datetime) -> int:
+    """Give up on jobs past their deadline; returns how many."""
+    from sqlalchemy import select
+    from .models import HostMaintenanceJob
+    gave_up = 0
+    for job in (await session.execute(select(HostMaintenanceJob).where(
+            HostMaintenanceJob.state.in_(ACTIVE)))).scalars():
+        if check_deadline(job, now):
+            gave_up += 1
+            await _window(session, job, await _host_name(session, job.host_id), now)
+    return gave_up
+
+
+class CannotStart(Exception):
+    """Why a host can't take a maintenance job now, in words for the page."""
+
+
+def pick_drainer(node_hostname_to_id: dict, managers: dict, host_name: str,
+                 host_id: str, helpers: set) -> tuple:
+    """(node_id, drain_host_id) for a host in a swarm, (None, None) for one
+    that isn't. `managers` maps a manager's host id to its swarm node id;
+    `helpers` is the hosts whose agent has the maintenance helper. The host
+    drains itself when it is a manager, else another manager does."""
+    node_id = node_hostname_to_id.get(host_name.lower())
+    if node_id is None:
+        return None, None
+    if host_id in managers and host_id in helpers:
+        return node_id, host_id
+    for mgr in sorted(managers):
+        if mgr != host_id and mgr in helpers:
+            return node_id, mgr
+    raise CannotStart("it is a swarm node and no swarm manager has the maintenance "
+                      "helper to drain it: run Update agent on a manager")
+
+
+async def plan_job(session, host, reg, reboot_policy: str = "if_required") -> tuple:
+    """(node_id, drain_host_id) for a job on `host`, or CannotStart saying why
+    it can't have one now. The panel asks this before offering the button."""
+    from sqlalchemy import select
+    from .models import HostAgentRegistration, HostMaintenanceJob
+    if reg is None or not reg.maint_helper:
+        raise CannotStart("its agent has no maintenance helper: run Update agent first")
+    if reboot_policy not in ("if_required", "never"):
+        raise CannotStart(f"unknown reboot policy {reboot_policy!r}")
+    busy = (await session.execute(select(HostMaintenanceJob.id).where(
+        HostMaintenanceJob.host_id == host.id,
+        HostMaintenanceJob.state.in_(ACTIVE)))).first()
+    if busy is not None:
+        raise CannotStart("a maintenance job is already running on it")
+    if reboot_policy == "never":
+        return None, None
+    nodes, managers = await _swarm(session)
+    helpers = set((await session.execute(select(HostAgentRegistration.host_id).where(
+        HostAgentRegistration.maint_helper.is_(True)))).scalars())
+    return pick_drainer(nodes, managers, host.name, host.id, helpers)
+
+
+async def create_job(session, host, reg, user_id: str | None,
+                     reboot_policy: str = "if_required", now: datetime | None = None):
+    """Queue a maintenance job for `host`; it starts with the agent's next post."""
+    from .models import HostMaintenanceJob
+    now = now or utcnow()
+    node_id, drain_host_id = await plan_job(session, host, reg, reboot_policy)
+    job = HostMaintenanceJob(host_id=host.id, state="queued", state_at=now,
+                             reboot_policy=reboot_policy, node_id=node_id,
+                             drain_host_id=drain_host_id, created_by=user_id, created_at=now)
+    session.add(job)
+    await session.flush()
+    return job
+
+
+async def latest_job(session, host_id: str, now: datetime, days: int = 7):
+    """The host's running job, else its last one in `days`, for the panel."""
+    from sqlalchemy import select
+    from .models import HostMaintenanceJob
+    return (await session.execute(select(HostMaintenanceJob).where(
+        HostMaintenanceJob.host_id == host_id,
+        HostMaintenanceJob.created_at >= now - timedelta(days=days),
+    ).order_by(HostMaintenanceJob.state.in_(ACTIVE).desc(),
+               HostMaintenanceJob.created_at.desc()).limit(1))).scalar_one_or_none()
+
+
+LABELS = {
+    "queued": "Waiting for the agent",
+    "updating": "Updating",
+    "draining": "Draining its swarm node",
+    "rebooting": "Rebooting",
+    "undraining": "Making its swarm node active again",
+    "done": "Done",
+    "failed": "Failed",
+    "timed_out": "Timed out",
+}
+
+
+def summary(job) -> str:
+    """One line on what the job did, for the panel."""
+    if job.state not in FINISHED:
+        return LABELS[job.state]
+    parts = []
+    if job.packages is not None:
+        parts.append(f"{job.packages} package{'s' if job.packages != 1 else ''} updated")
+    if job.back_at is not None:
+        parts.append("rebooted")
+    elif job.reboot_needed and job.reboot_policy == "never":
+        parts.append("reboot needed, not allowed by this job")
+    elif job.state == "done" and job.reboot_needed is False:
+        parts.append("no reboot needed")
+    head = LABELS[job.state]
+    return f"{head}: {', '.join(parts)}" if parts else head
+
+
+async def _swarm(session) -> tuple:
+    """({node hostname: node id}, {manager host id: its node id}) from what the
+    Docker plugin last stored; empty when it isn't installed. A manager is a
+    host whose agent reported the swarm's nodes."""
+    try:
+        from plugins.docker.models import DockerSwarmNode
+    except ImportError:
+        return {}, {}
+    from sqlalchemy import select
+    from steward.models.hosts import Host
+    rows = (await session.execute(select(DockerSwarmNode))).scalars().all()
+    nodes = {(n.hostname or "").lower(): n.node_id for n in rows if n.hostname}
+    reporters = {n.host_id for n in rows}
+    names = dict((await session.execute(select(Host.id, Host.name).where(
+        Host.id.in_(reporters)))).all()) if reporters else {}
+    managers = {hid: nodes.get(name.lower()) for hid, name in names.items()}
+    return nodes, managers