fix(cluster): rebase onto dev + address Kilo findings (1C/1W/2S) for worker drain

Rebased feat/worker-self-drain onto origin/dev, resolving conflicts in
manager.py, routes/cluster.py, and worker/agent.py from recently-landed
registration-drift refresh fields (#1538).

Kilo fixes:
- CRITICAL: Protect 'update-available' from status-less heartbeat reversion.
  The periodic ~15s heartbeat with no status was silently flipping
  update-available workers back to online, defeating the update signal.
  Now 'update-available' is protected alongside 'draining' — only explicit
  status transitions are honoured.
- WARNING: Handle 'updating' in _monitor_loop. A worker stuck/crashed in
  'updating' state was never re-onlined or offlined, permanently excluded
  from routing/catalog. Now they are force-offlined after HEARTBEAT_TIMEOUT.
- SUGGESTION: Validate status against allow-list and sanitize drain_reason
  in notification block to prevent operator-facing UI injection.
- SUGGESTION: notify_drain_complete() now sends status='updating' heartbeat
  so the controller transitions from 'draining' to 'updating'.

Test update: test_heartbeat_without_status_reonlines_update_available
renamed to test_heartbeat_without_status_preserves_update_available and
now asserts the corrected protection behavior.

72/72 targeted tests pass.

Author: hognek

tests/test_cluster.py

@@ -520,19 +520,29 @@ async def test_drain_preserved_on_subsequent_heartbeats(self):
         assert mgr.get_worker("gpu-box").status == "draining"
         assert mgr.get_worker("gpu-box").load == 0.5
 
-    async def test_heartbeat_without_status_reonlines_update_available(self):
-        """A normal heartbeat without status should transition
-        update-available back to online (it's not a drain)."""
+    async def test_heartbeat_without_status_preserves_update_available(self):
+        """A status-less heartbeat preserves update-available (taOS #890 C2).
+
+        The periodic ~15s heartbeat carries no status; when combined with the
+        old manager.py:226-231 logic that flipped any non-draining worker back
+        to "online", this silently cancelled the update signal before the
+        worker could initiate its drain.  Now "update-available" is protected
+        in the same way "draining" is: only an explicit status transition
+        (e.g. "draining" or "updating") changes it."""
         mgr = ClusterManager()
         await mgr.register_worker(_make_worker("gpu-box", capabilities=["chat"]))
 
         mgr.heartbeat("gpu-box", status="update-available")
         assert mgr.get_worker("gpu-box").status == "update-available"
 
-        # Normal heartbeat without explicit status
+        # Normal heartbeat without explicit status — should stay update-available
         mgr.heartbeat("gpu-box", load=0.2)
-        # Should go back to online because update-available is not draining
-        assert mgr.get_worker("gpu-box").status == "online"
+        assert mgr.get_worker("gpu-box").status == "update-available"
+        assert mgr.get_worker("gpu-box").load == 0.2
+
+        # Explicit transition to draining still works
+        mgr.heartbeat("gpu-box", status="draining", drain_reason="update")
+        assert mgr.get_worker("gpu-box").status == "draining"
 
     async def test_monitor_loop_offlines_stale_update_available_worker(self):
         """An update-available worker that stops heartbeating is marked offline."""

tinyagentos/cluster/manager.py

@@ -231,9 +231,12 @@ def heartbeat(
         # A worker mid graceful-drain keeps heartbeating while it finishes
         # inflight work. Do NOT flip it back to "online" — that would re-enter
         # it into routing/catalog/lease-claims before its leases drain and
-        # defeat the drain (taOS #890). Leave the drain to complete or be
-        # cancelled explicitly; every other status re-onlines as before.
-        if worker.status != "draining":
+        # defeat the drain (taOS #890). Similarly, protect "update-available"
+        # so a periodic status-less heartbeat doesn't silently cancel the
+        # update signal before the worker can initiate the drain.
+        # Only block status-less heartbeats; explicit status transitions
+        # (e.g. "update-available" → "draining") are still honoured.
+        if worker.status not in ("draining", "update-available") or status is not None:
             # Honour worker-initiated status transitions (taOS #890 C2).
             if status in ("draining", "update-available", "updating"):
                 worker.status = status
@@ -306,8 +309,11 @@ def heartbeat(
         # Worker-initiated drain notification (taOS #890 C2).
         # Emit when the worker transitions into draining/update-available on
         # its own initiative, so the operator sees it in the activity feed.
-        if self._notifications and status in ("draining", "update-available", "updating") and prev_status not in (status,):
-            reason = drain_reason or "unspecified"
+        # Validate: only trusted status values; sanitize drain_reason to
+        # prevent injection into notification UI.
+        _VALID_STATUSES = frozenset({"draining", "update-available", "updating"})
+        if self._notifications and status in _VALID_STATUSES and prev_status not in (status,):
+            reason = (drain_reason or "unspecified").replace("'", "\\'").replace("\\", "\\\\")[:200]
             event_type = f"worker.{status}" if status != "draining" else "worker.drain"
             title_map = {
                 "draining": f"Worker '{worker.name}' self-initiated drain",
@@ -824,6 +830,30 @@ async def _monitor_loop(self):
                                 logger.exception(
                                     "gpu-arbiter: cancel for stale-drain of '%s' failed",
                                     worker.name)
+                # Handle updating workers: if they stop heartbeating they're
+                # stuck in a bad state — force-offline them so they can re-onboard.
+                elif worker.status == "updating":
+                    if (now - worker.last_heartbeat) > HEARTBEAT_TIMEOUT:
+                        worker.status = "offline"
+                        logger.warning(
+                            "Worker '%s' stuck in 'updating' (no heartbeat for %ds) — marked offline",
+                            worker.name, HEARTBEAT_TIMEOUT,
+                        )
+                        if self._notifications:
+                            await self._notifications.emit_event(
+                                "worker.leave",
+                                f"Worker '{worker.name}' timed out during update",
+                                f"Worker was stuck in 'updating' state. Forced offline to allow re-onboarding.",
+                                level="warning",
+                            )
+                        async with self._lease_lock:
+                            update_lids = [
+                                lid for lid, lease in self._leases.items()
+                                if (parsed := self._parse_resource_id(lease.resource_id))
+                                and parsed[0] == worker.name
+                            ]
+                            for lid in update_lids:
+                                self._leases.pop(lid, None)
             async with self._lease_lock:
                 self._sweep_expired_leases()
             await asyncio.sleep(5)

tinyagentos/worker/agent.py

@@ -717,13 +717,14 @@ async def notify_drain_complete(self) -> int:
         """Notify the controller that the worker's drain is complete.
 
         After in-flight work finishes, the worker sends one final
-        heartbeat to confirm readiness for the update. The controller
-        can then proceed with the update deploy.
+        heartbeat with status="updating" to signal readiness for
+        the update. The controller can then proceed with the update
+        deploy.
 
         Returns the HTTP status code from the controller.
         """
         logger.info("worker '%s': drain complete, ready for update", self.name)
-        return await self.heartbeat()
+        return await self.heartbeat(status="updating")
 
     def _log_repair_instruction(self) -> None:
         logger.error(