From c6d9a985e79ac994bbe99360c286753d5789feb8 Mon Sep 17 00:00:00 2001
From: Hogne <227774406+hognek@users.noreply.github.com>
Date: Mon, 6 Jul 2026 23:17:52 +0200
Subject: [PATCH 1/4] feat(cluster): worker auto-update with graceful drain,
 pause, install, restart
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Adds drain_worker/cancel_drain to ClusterManager for graceful worker
detach without dropping inflight GPU tasks (taOS #890).

ClusterManager:
- drain_worker(name, graceful=True): worker enters 'draining' status,
  excluded from routing/catalog/lease claims. When graceful=False,
  all leases are force-released and worker is marked offline.
- cancel_drain(name): returns draining worker to 'online'.
- _monitor_loop: auto-completes drain when all leases released;
  force-finishes stale drains on heartbeat timeout.
- _worker_for_resource, get_workers_for_capability, aggregate_catalog:
  all exclude draining workers.

Routes:
- POST /api/cluster/workers/{name}/drain — begin graceful/force drain
- POST /api/cluster/workers/{name}/cancel-drain — cancel in-progress drain
- POST /api/cluster/workers/{name}/update — full auto-update orchestration:
  drain → deploy update-worker → restart → re-register

Tests: 9 new tests (25 total in test_cluster.py), 191 cluster tests pass.
---
 tests/test_cluster.py          | 118 ++++++++++++++++++++++++++
 tinyagentos/cluster/manager.py | 150 ++++++++++++++++++++++++++++++---
 tinyagentos/routes/cluster.py  |  93 ++++++++++++++++++++
 3 files changed, 351 insertions(+), 10 deletions(-)

diff --git a/tests/test_cluster.py b/tests/test_cluster.py
index 74875d441..111b843b5 100644
--- a/tests/test_cluster.py
+++ b/tests/test_cluster.py
@@ -259,3 +259,121 @@ async def test_router_returns_none_for_no_workers(self):
         data, worker_name = await router.route_request("chat", "POST", "/v1/chat/completions", {})
         assert data is None
         assert worker_name is None
+
+
+@pytest.mark.asyncio
+class TestWorkerDrain:
+    """Tests for taOS #890 — worker auto-update with graceful drain."""
+
+    async def test_drain_worker_graceful_sets_draining_status(self):
+        """Graceful drain sets status to 'draining', leaves leases intact."""
+        mgr = ClusterManager()
+        await mgr.register_worker(_make_worker("gpu-box", capabilities=["chat"]))
+
+        result = mgr.drain_worker("gpu-box", graceful=True)
+        assert result["worker"] == "gpu-box"
+        assert result["previous_status"] == "online"
+        assert result["status"] == "draining"
+        assert result["released_leases"] == 0
+        assert mgr.get_worker("gpu-box").status == "draining"
+
+    async def test_drain_worker_force_releases_leases_and_marks_offline(self):
+        """Force drain releases all leases and marks worker offline."""
+        mgr = ClusterManager()
+        await mgr.register_worker(_make_worker("gpu-box", capabilities=["chat"]))
+
+        # Add a lease for this worker
+        from tinyagentos.cluster.worker_protocol import GpuLease
+        import time
+        lease = GpuLease(
+            lease_id="l_test1",
+            resource_id="gpu-box:gpu-cuda-0",
+            caller="test",
+            expires_at=time.time() + 60,
+            required_vram_mb=0,
+        )
+        mgr._leases["l_test1"] = lease
+
+        result = mgr.drain_worker("gpu-box", graceful=False)
+        assert result["released_leases"] == 1
+        assert result["status"] == "offline"
+        assert mgr.get_worker("gpu-box").status == "offline"
+        assert "l_test1" not in mgr._leases
+
+    async def test_drain_worker_unknown_returns_error(self):
+        """Draining a non-existent worker returns error dict."""
+        mgr = ClusterManager()
+        result = mgr.drain_worker("nonexistent")
+        assert "error" in result
+        assert result["worker"] == "nonexistent"
+
+    async def test_cancel_drain_returns_worker_to_online(self):
+        """Cancel drain returns a draining worker back to online."""
+        mgr = ClusterManager()
+        await mgr.register_worker(_make_worker("gpu-box", capabilities=["chat"]))
+        mgr.drain_worker("gpu-box", graceful=True)
+        assert mgr.get_worker("gpu-box").status == "draining"
+
+        result = mgr.cancel_drain("gpu-box")
+        assert result["worker"] == "gpu-box"
+        assert result["status"] == "online"
+        assert mgr.get_worker("gpu-box").status == "online"
+
+    async def test_cancel_drain_only_works_on_draining(self):
+        """Cancel drain returns error for non-draining workers."""
+        mgr = ClusterManager()
+        await mgr.register_worker(_make_worker("gpu-box"))
+
+        result = mgr.cancel_drain("gpu-box")
+        assert "error" in result
+        assert "not draining" in result["error"]
+
+    async def test_cancel_drain_unknown_returns_error(self):
+        """Cancel drain on unknown worker returns error."""
+        mgr = ClusterManager()
+        result = mgr.cancel_drain("nonexistent")
+        assert "error" in result
+
+    async def test_draining_workers_excluded_from_routing(self):
+        """Draining workers should not receive new tasks."""
+        mgr = ClusterManager()
+        await mgr.register_worker(_make_worker("online-gpu", capabilities=["chat"], load=0.1))
+        await mgr.register_worker(_make_worker("draining-gpu", capabilities=["chat"], load=0.0))
+        mgr.drain_worker("draining-gpu", graceful=True)
+
+        result = mgr.get_workers_for_capability("chat")
+        assert len(result) == 1
+        assert result[0].name == "online-gpu"
+
+    async def test_draining_workers_excluded_from_catalog(self):
+        """Draining workers are excluded from the aggregate catalog."""
+        mgr = ClusterManager()
+        online = _make_worker("online", capabilities=["chat"])
+        online.backends = [{"name": "b1", "type": "ollama", "url": "u", "capabilities": ["chat"], "models": [{"name": "m1"}]}]
+        draining = _make_worker("draining", capabilities=["chat"])
+        draining.backends = [{"name": "b2", "type": "ollama", "url": "u", "capabilities": ["chat"], "models": [{"name": "m2"}]}]
+        await mgr.register_worker(online)
+        await mgr.register_worker(draining)
+        mgr.drain_worker("draining", graceful=True)
+
+        out = mgr.aggregate_catalog()
+        assert [w["name"] for w in out["workers"]] == ["online"]
+        assert [m["name"] for m in out["models"]] == ["m1"]
+
+    async def test_draining_workers_excluded_from_lease_claim(self):
+        """Lease claims should fail for draining workers."""
+        mgr = ClusterManager()
+        await mgr.register_worker(_make_worker("draining-gpu", url="http://draining:9000"))
+
+        # Set up worker with VRAM info
+        worker = mgr.get_worker("draining-gpu")
+        worker.free_vram_mb = 8000
+
+        mgr.drain_worker("draining-gpu", graceful=True)
+
+        lease = await mgr.claim_lease(
+            resource_id="draining-gpu:gpu-cuda-0",
+            caller="test",
+            ttl_seconds=30,
+        )
+        assert lease is None
diff --git a/tinyagentos/cluster/manager.py b/tinyagentos/cluster/manager.py
index 69df08323..419e4111b 100644
--- a/tinyagentos/cluster/manager.py
+++ b/tinyagentos/cluster/manager.py
@@ -273,13 +273,15 @@ def _parse_resource_id(self, resource_id: str) -> tuple[str, str] | None:
         return worker, rest
 
     def _worker_for_resource(self, resource_id: str) -> WorkerInfo | None:
-        """Return the WorkerInfo for a resource_id, or None."""
+        """Return the WorkerInfo for a resource_id, or None.
+
+        Excludes draining and offline workers (taOS #890)."""
         parsed = self._parse_resource_id(resource_id)
         if parsed is None:
             return None
         worker_name, _ = parsed
         worker = self._workers.get(worker_name)
-        if worker is None or worker.status != "online":
+        if worker is None or worker.status not in ("online",):
             return None
         return worker
 
@@ -393,10 +395,106 @@ def _sweep_expired_leases(self):
             lease = self._leases.pop(lid)
             logger.debug("Lease expired: %s on %s", lid, lease.resource_id)
 
+    # ── taOS #890: graceful worker drain for auto-update ───────────────
+
+    def drain_worker(self, name: str, graceful: bool = True) -> dict:
+        """Begin draining a worker — gracefully detach without dropping tasks.
+
+        When ``graceful=True`` (the default):
+        1. Worker enters ``"draining"`` status — no new tasks routed to it.
+        2. Existing leases are allowed to run to completion (TTL not touched).
+        3. The monitor loop will eventually sweep expired leases and then
+           mark the worker ``"offline"`` once all leases are released.
+
+        When ``graceful=False``:
+        1. Worker enters ``"draining"`` status.
+        2. All active leases for this worker are released immediately.
+        3. Worker is marked ``"offline"`` immediately.
+
+        Returns a dict with ``worker``, ``previous_status``, ``released_leases``,
+        and ``status``.
+        """
+        worker = self._workers.get(name)
+        if worker is None:
+            return {"worker": name, "error": "worker not found"}
+        prev_status = worker.status
+        worker.status = "draining"
+        logger.info("Worker '%s' entering drain (was %s, graceful=%s)",
+                     name, prev_status, graceful)
+
+        released = 0
+
+        if not graceful:
+            # Force-release all leases for this worker
+            lids = [
+                lid for lid, lease in self._leases.items()
+                if lease.resource_id.startswith(name + ":")
+            ]
+            for lid in lids:
+                self._leases.pop(lid, None)
+                released += 1
+            logger.info("Worker '%s' drain: force-released %d leases", name, released)
+            worker.status = "offline"
+
+        if self._notifications:
+            detail = (
+                f"Worker '{name}' is draining. "
+                f"{'Tasks will complete before detach.' if graceful else 'All leases released immediately.'}"
+            )
+            try:
+                asyncio.get_running_loop().create_task(
+                    self._notifications.emit_event(
+                        "worker.drain",
+                        f"Worker '{name}' draining",
+                        detail,
+                        level="info",
+                    )
+                )
+            except RuntimeError:
+                pass
+
+        return {
+            "worker": name,
+            "previous_status": prev_status,
+            "status": worker.status,
+            "released_leases": released,
+        }
+
+    def cancel_drain(self, name: str) -> dict:
+        """Cancel an in-progress drain and return the worker to online.
+
+        Only works when the worker is in ``"draining"`` status and has
+        not yet been marked offline.
+        """
+        worker = self._workers.get(name)
+        if worker is None:
+            return {"worker": name, "error": "worker not found"}
+        if worker.status != "draining":
+            return {"worker": name, "error": f"worker is not draining (status={worker.status})"}
+        worker.status = "online"
+        logger.info("Worker '%s' drain cancelled — back online", name)
+
+        if self._notifications:
+            try:
+                asyncio.get_running_loop().create_task(
+                    self._notifications.emit_event(
+                        "worker.online",
+                        f"Worker '{name}' drain cancelled",
+                        f"Worker '{name}' is back online after drain was cancelled.",
+                        level="info",
+                    )
+                )
+            except RuntimeError:
+                pass
+
+        return {"worker": name, "status": "online"}
+
     # ── Capability routing ─────────────────────────────────────────────
 
     def get_workers_for_capability(self, capability: str) -> list[WorkerInfo]:
-        """Get online workers that support a capability, sorted by priority (lowest load first)."""
+        """Get online workers that support a capability, sorted by priority (lowest load first).
+
+        Draining workers are excluded (taOS #890)."""
         eligible = [
             w for w in self._workers.values()
             if w.status == "online" and capability in w.capabilities
@@ -443,7 +541,7 @@ def aggregate_catalog(self) -> dict:
 
         for worker in self._workers.values():
             if worker.status != "online":
-                continue
+                continue  # skip offline and draining workers (taOS #890)
 
             worker_caps = set(worker.capabilities or [])
             all_capabilities |= worker_caps
@@ -486,14 +584,12 @@ def aggregate_catalog(self) -> dict:
 
     async def _monitor_loop(self):
         """Monitor worker heartbeats, mark stale workers as offline, and
-        sweep expired GPU leases."""
+        sweep expired GPU leases.  Auto-completes draining workers
+        whose leases have all been released (taOS #890)."""
         while True:
             now = time.time()
-            for worker in self._workers.values():
-                # The 'local' worker is the controller itself — it never sends
-                # heartbeats (it IS the server), so never mark it offline.
-                if worker.name == "local":
-                    continue
+            for worker in list(self._workers.values()):
+                # Handle online workers that haven't heartbeated
                 if worker.status == "online" and (now - worker.last_heartbeat) > HEARTBEAT_TIMEOUT:
                     worker.status = "offline"
                     logger.warning(f"Worker '{worker.name}' marked offline (no heartbeat for {HEARTBEAT_TIMEOUT}s)")
@@ -516,6 +612,40 @@ async def _monitor_loop(self):
                         for lid in offline_lids:
                             self._leases.pop(lid, None)
                             logger.debug("Lease %s released — worker %s went offline", lid, worker.name)
+
+                # Handle draining workers: auto-complete if no active leases (taOS #890)
+                elif worker.status == "draining":
+                    active_leases = [
+                        lid for lid, lease in self._leases.items()
+                        if lease.resource_id.startswith(worker.name + ":")
+                    ]
+                    if not active_leases:
+                        worker.status = "offline"
+                        logger.info(
+                            "Worker '%s' drained — all leases released, marked offline",
+                            worker.name,
+                        )
+                        if self._notifications:
+                            await self._notifications.emit_event(
+                                "worker.leave",
+                                f"Worker '{worker.name}' drained and went offline",
+                                "All tasks completed; worker detached gracefully.",
+                                level="info",
+                            )
+                    elif (now - worker.last_heartbeat) > HEARTBEAT_TIMEOUT:
+                        # Draining worker went stale — force-finish the drain
+                        lids = [
+                            lid for lid, lease in self._leases.items()
+                            if lease.resource_id.startswith(worker.name + ":")
+                        ]
+                        for lid in lids:
+                            self._leases.pop(lid, None)
+                        worker.status = "offline"
+                        logger.warning(
+                            "Worker '%s' drain timed out (no heartbeat for %ds) — "
+                            "force-released %d leases, marked offline",
+                            worker.name, HEARTBEAT_TIMEOUT, len(lids),
+                        )
             async with self._lease_lock:
                 self._sweep_expired_leases()
             await asyncio.sleep(5)
diff --git a/tinyagentos/routes/cluster.py b/tinyagentos/routes/cluster.py
index ef5d91708..338d82284 100644
--- a/tinyagentos/routes/cluster.py
+++ b/tinyagentos/routes/cluster.py
@@ -1101,3 +1101,96 @@ async def list_leases(request: Request):
         ],
         "count": len(leases),
     }
+
+
+# ── taOS #890: worker auto-update with graceful drain ────────────────────
+
+
+class DrainRequest(BaseModel):
+    graceful: bool = True
+
+
+@router.post("/api/cluster/workers/{name}/drain")
+async def drain_worker(request: Request, name: str, body: DrainRequest = DrainRequest()):
+    """Begin draining a worker — gracefully detach without dropping tasks.
+
+    When ``graceful=true`` (the default), the worker enters "draining" status:
+    no new tasks are routed to it, but existing leases run to completion.
+    The monitor loop auto-completes the drain when all leases are released.
+
+    When ``graceful=false``, all leases are force-released and the worker
+    is marked offline immediately.
+    """
+    cluster = request.app.state.cluster_manager
+    result = cluster.drain_worker(name, graceful=body.graceful)
+    if "error" in result:
+        return JSONResponse(result, status_code=404)
+    return result
+
+
+@router.post("/api/cluster/workers/{name}/cancel-drain")
+async def cancel_drain(request: Request, name: str):
+    """Cancel an in-progress drain and return the worker to online status."""
+    cluster = request.app.state.cluster_manager
+    result = cluster.cancel_drain(name)
+    if "error" in result:
+        return JSONResponse(result, status_code=404)
+    return result
+
+
+@router.post("/api/cluster/workers/{name}/update")
+async def update_worker(request: Request, name: str):
+    """Trigger a full auto-update sequence on a worker.
+
+    Orchestrates the graceful update sequence:
+    1. Pause incoming tasks — worker enters "draining" status.
+    2. Inflight leases run to completion (drain).
+    3. Send ``update-worker`` deploy command to the worker.
+    4. Worker restarts, re-registers, returns to "online".
+    5. Routing resumes automatically once the worker is back online.
+    """
+    cluster = request.app.state.cluster_manager
+    worker = cluster.get_worker(name)
+    if not worker:
+        return JSONResponse({"error": f"Worker '{name}' not found"}, status_code=404)
+    if worker.status not in ("online", "draining"):
+        return JSONResponse(
+            {"error": f"Worker '{name}' is not online (status={worker.status})"},
+            status_code=400,
+        )
+
+    # Step 1: Begin draining
+    drain_result = cluster.drain_worker(name, graceful=True)
+    if "error" in drain_result:
+        return JSONResponse(drain_result, status_code=404)
+
+    # Step 2: Trigger the update-worker deploy command on the worker
+    import httpx
+    deploy_result = None
+    try:
+        async with httpx.AsyncClient(timeout=620) as client:
+            resp = await client.post(
+                f"{worker.url}/api/worker/deploy",
+                json={"command": "update-worker"},
+            )
+            deploy_result = resp.json()
+    except Exception as exc:
+        # Deploy failed — cancel drain so worker can still serve traffic
+        cluster.cancel_drain(name)
+        return JSONResponse(
+            {"error": f"Worker update deploy failed: {exc}", "drain_cancelled": True},
+            status_code=502,
+        )
+
+    return {
+        "worker": name,
+        "status": "updating",
+        "previous_status": drain_result["previous_status"],
+        "drain": drain_result,
+        "deploy": deploy_result,
+        "message": (
+            "Worker is draining and updating. It will restart and re-register. "
+            "The controller monitor loop will auto-complete the drain once "
+            "all leases are released and the new worker process heartbeats."
+        ),
+    }

From 6903c0032f3e715e948e7975478b1748d9433e5b Mon Sep 17 00:00:00 2001
From: Hogne <227774406+hognek@users.noreply.github.com>
Date: Tue, 7 Jul 2026 13:39:09 +0200
Subject: [PATCH 2/4] fix(cluster): use exact worker-name compare for drain
 lease matching

Replace startswith(name + ':') with _parse_resource_id + exact
parsed[0] == name compare in three drain code paths to prevent
accidental cross-worker lease release when worker names share
a prefix (e.g. 'worker-a' matching 'worker-a-pro' leases).

Fixes #1690
---
 tinyagentos/cluster/manager.py | 9 ++++++---
 1 file changed, 6 insertions(+), 3 deletions(-)

diff --git a/tinyagentos/cluster/manager.py b/tinyagentos/cluster/manager.py
index 419e4111b..f7025de7b 100644
--- a/tinyagentos/cluster/manager.py
+++ b/tinyagentos/cluster/manager.py
@@ -428,7 +428,8 @@ def drain_worker(self, name: str, graceful: bool = True) -> dict:
             # Force-release all leases for this worker
             lids = [
                 lid for lid, lease in self._leases.items()
-                if lease.resource_id.startswith(name + ":")
+                if (parsed := self._parse_resource_id(lease.resource_id))
+                and parsed[0] == name
             ]
             for lid in lids:
                 self._leases.pop(lid, None)
@@ -617,7 +618,8 @@ async def _monitor_loop(self):
                 elif worker.status == "draining":
                     active_leases = [
                         lid for lid, lease in self._leases.items()
-                        if lease.resource_id.startswith(worker.name + ":")
+                        if (parsed := self._parse_resource_id(lease.resource_id))
+                        and parsed[0] == worker.name
                     ]
                     if not active_leases:
                         worker.status = "offline"
@@ -636,7 +638,8 @@ async def _monitor_loop(self):
                         # Draining worker went stale — force-finish the drain
                         lids = [
                             lid for lid, lease in self._leases.items()
-                            if lease.resource_id.startswith(worker.name + ":")
+                            if (parsed := self._parse_resource_id(lease.resource_id))
+                            and parsed[0] == worker.name
                         ]
                         for lid in lids:
                             self._leases.pop(lid, None)

From a46b09b930efc7245a801f28bcbbfdb4a4472a25 Mon Sep 17 00:00:00 2001
From: Hogne <227774406+hognek@users.noreply.github.com>
Date: Fri, 10 Jul 2026 22:00:59 +0200
Subject: [PATCH 3/4] fix(cluster): address review blockers on worker
 auto-update PR #1690

- Restore local worker guard in _monitor_loop (prevents controller
  from marking itself offline and dropping its own leases).
- Fix update_worker error handling: use raise_for_status() for real
  HTTP errors; catch expected disconnect during worker restart
  without cancelling drain; keep drain active until re-registration.
- Move force-path lease pops under _lease_lock for consistency with
  the online->offline path.
- Add idempotency guard: /update returns 409 if worker is already
  draining.
- Move httpx import to module level in routes/cluster.py (matches
  codebase convention).
- Add 7 tests for update_worker covering all error paths.
---
 tests/test_routes_cluster.py   | 163 +++++++++++++++++++++++++++++++++
 tinyagentos/cluster/manager.py |  25 +++--
 tinyagentos/routes/cluster.py  |  42 +++++++--
 3 files changed, 214 insertions(+), 16 deletions(-)

diff --git a/tests/test_routes_cluster.py b/tests/test_routes_cluster.py
index 8daf0848d..d804d2b43 100644
--- a/tests/test_routes_cluster.py
+++ b/tests/test_routes_cluster.py
@@ -520,3 +520,166 @@ async def fake_remote_list():
     assert fedora is not None
     assert fedora["hardware_known"] is True, fedora
     assert fedora["tier_id"] not in ("", "unknown"), fedora
+
+
+# ── update_worker endpoint tests (taOS #890) ──────────────────────────────
+
+
+@pytest.mark.asyncio
+async def test_update_worker_not_found(client):
+    """404 when worker does not exist."""
+    resp = await client.post("/api/cluster/workers/nonexistent/update")
+    assert resp.status_code == 404
+    assert "not found" in resp.json()["error"]
+
+
+@pytest.mark.asyncio
+async def test_update_worker_already_draining(client, app):
+    """409 Conflict when worker is already draining (idempotency guard)."""
+    from tinyagentos.cluster.worker_protocol import WorkerInfo
+    cluster = app.state.cluster_manager
+    cluster._workers["gpu-box"] = WorkerInfo(  # noqa: SLF001
+        name="gpu-box", url="http://gpu:9000", status="draining",
+    )
+    resp = await client.post("/api/cluster/workers/gpu-box/update")
+    assert resp.status_code == 409
+    assert "already updating" in resp.json()["error"]
+
+
+@pytest.mark.asyncio
+async def test_update_worker_offline(client, app):
+    """400 when worker is offline."""
+    from tinyagentos.cluster.worker_protocol import WorkerInfo
+    cluster = app.state.cluster_manager
+    cluster._workers["gpu-box"] = WorkerInfo(  # noqa: SLF001
+        name="gpu-box", url="http://gpu:9000", status="offline",
+    )
+    resp = await client.post("/api/cluster/workers/gpu-box/update")
+    assert resp.status_code == 400
+    assert "not online" in resp.json()["error"]
+
+
+@pytest.mark.asyncio
+async def test_update_worker_success(client, app):
+    """Happy path: drain succeeds, deploy returns 200."""
+    from tinyagentos.cluster.worker_protocol import WorkerInfo
+    cluster = app.state.cluster_manager
+    cluster._workers["gpu-box"] = WorkerInfo(  # noqa: SLF001
+        name="gpu-box", url="http://gpu:9000", status="online",
+    )
+
+    from unittest.mock import MagicMock
+
+    mock_resp = MagicMock()
+    mock_resp.status_code = 200
+    mock_resp.json.return_value = {"ok": True, "command": "update-worker"}
+    mock_resp.raise_for_status = lambda: None
+
+    mock_client = AsyncMock()
+    mock_client.__aenter__.return_value = mock_client
+    mock_client.__aexit__.return_value = None
+    mock_client.post.return_value = mock_resp
+
+    with patch("tinyagentos.routes.cluster.httpx.AsyncClient", return_value=mock_client):
+        resp = await client.post("/api/cluster/workers/gpu-box/update")
+
+    assert resp.status_code == 200
+    data = resp.json()
+    assert data["worker"] == "gpu-box"
+    assert data["status"] == "updating"
+    assert data["drain"]["status"] == "draining"
+    assert data["deploy"]["ok"] is True
+
+    # Worker should be draining after the call
+    assert cluster.get_worker("gpu-box").status == "draining"
+
+
+@pytest.mark.asyncio
+async def test_update_worker_deploy_http_error(client, app):
+    """502 when deploy endpoint returns a non-2xx status — drain cancelled."""
+    from tinyagentos.cluster.worker_protocol import WorkerInfo
+    from unittest.mock import MagicMock
+    import httpx as _httpx
+    cluster = app.state.cluster_manager
+    cluster._workers["gpu-box"] = WorkerInfo(  # noqa: SLF001
+        name="gpu-box", url="http://gpu:9000", status="online",
+    )
+
+    mock_resp = MagicMock()
+    mock_resp.status_code = 500
+    mock_resp.raise_for_status = lambda: (_ for _ in ()).throw(
+        _httpx.HTTPStatusError("Server error", request=MagicMock(), response=mock_resp)
+    )
+
+    mock_client = AsyncMock()
+    mock_client.__aenter__.return_value = mock_client
+    mock_client.__aexit__.return_value = None
+    mock_client.post.return_value = mock_resp
+
+    with patch("tinyagentos.routes.cluster.httpx.AsyncClient", return_value=mock_client):
+        resp = await client.post("/api/cluster/workers/gpu-box/update")
+
+    assert resp.status_code == 502
+    data = resp.json()
+    assert "rejected by worker" in data["error"]
+    assert data["drain_cancelled"] is True
+
+    # Drain should be cancelled — worker back online
+    assert cluster.get_worker("gpu-box").status == "online"
+
+
+@pytest.mark.asyncio
+async def test_update_worker_deploy_disconnect_during_restart(client, app):
+    """200 when worker disconnects mid-update (expected restart).  Drain
+    stays active — the worker will re-register and the monitor loop
+    will auto-complete the drain once all leases are released."""
+    from tinyagentos.cluster.worker_protocol import WorkerInfo
+    import httpx as _httpx
+    cluster = app.state.cluster_manager
+    cluster._workers["gpu-box"] = WorkerInfo(  # noqa: SLF001
+        name="gpu-box", url="http://gpu:9000", status="online",
+    )
+
+    mock_client = AsyncMock()
+    mock_client.__aenter__.return_value = mock_client
+    mock_client.__aexit__.return_value = None
+    mock_client.post.side_effect = _httpx.RemoteProtocolError(
+        "Server disconnected", request=AsyncMock()
+    )
+
+    with patch("tinyagentos.routes.cluster.httpx.AsyncClient", return_value=mock_client):
+        resp = await client.post("/api/cluster/workers/gpu-box/update")
+
+    assert resp.status_code == 200
+    data = resp.json()
+    assert data["worker"] == "gpu-box"
+    assert data["status"] == "updating"
+    assert data["deploy"]["status"] == "acknowledged"
+    assert "restart expected" in data["deploy"]["note"]
+
+    # Drain must NOT be cancelled — worker should still be draining
+    assert cluster.get_worker("gpu-box").status == "draining"
+
+
+@pytest.mark.asyncio
+async def test_update_worker_deploy_unexpected_error(client, app):
+    """502 on unexpected exception — drain cancelled."""
+    from tinyagentos.cluster.worker_protocol import WorkerInfo
+    cluster = app.state.cluster_manager
+    cluster._workers["gpu-box"] = WorkerInfo(  # noqa: SLF001
+        name="gpu-box", url="http://gpu:9000", status="online",
+    )
+
+    mock_client = AsyncMock()
+    mock_client.__aenter__.return_value = mock_client
+    mock_client.__aexit__.return_value = None
+    mock_client.post.side_effect = OSError("Network unreachable")
+
+    with patch("tinyagentos.routes.cluster.httpx.AsyncClient", return_value=mock_client):
+        resp = await client.post("/api/cluster/workers/gpu-box/update")
+
+    assert resp.status_code == 502
+    data = resp.json()
+    assert "deploy failed" in data["error"]
+    assert data["drain_cancelled"] is True
+    assert cluster.get_worker("gpu-box").status == "online"
diff --git a/tinyagentos/cluster/manager.py b/tinyagentos/cluster/manager.py
index f7025de7b..9a7ee8b08 100644
--- a/tinyagentos/cluster/manager.py
+++ b/tinyagentos/cluster/manager.py
@@ -590,6 +590,11 @@ async def _monitor_loop(self):
         while True:
             now = time.time()
             for worker in list(self._workers.values()):
+                # The local worker (controller itself) must never be marked
+                # offline by the monitor loop — doing so would drop the
+                # controller's own leases and break local routing.
+                if worker.name == "local":
+                    continue
                 # Handle online workers that haven't heartbeated
                 if worker.status == "online" and (now - worker.last_heartbeat) > HEARTBEAT_TIMEOUT:
                     worker.status = "offline"
@@ -635,14 +640,18 @@ async def _monitor_loop(self):
                                 level="info",
                             )
                     elif (now - worker.last_heartbeat) > HEARTBEAT_TIMEOUT:
-                        # Draining worker went stale — force-finish the drain
-                        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 lids:
-                            self._leases.pop(lid, None)
+                        # Draining worker went stale — force-finish the drain.
+                        # Lease pops must be serialised under _lease_lock so
+                        # concurrent claim/release operations see a consistent
+                        # view of the lease table.
+                        async with self._lease_lock:
+                            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 lids:
+                                self._leases.pop(lid, None)
                         worker.status = "offline"
                         logger.warning(
                             "Worker '%s' drain timed out (no heartbeat for %ds) — "
diff --git a/tinyagentos/routes/cluster.py b/tinyagentos/routes/cluster.py
index 338d82284..8db2df984 100644
--- a/tinyagentos/routes/cluster.py
+++ b/tinyagentos/routes/cluster.py
@@ -2,17 +2,17 @@
 
 import asyncio
 import logging
+import re
 import subprocess
 from dataclasses import asdict
 
+import httpx
 from fastapi import APIRouter, Request
 from fastapi.responses import JSONResponse
 from pydantic import BaseModel
 
 logger = logging.getLogger(__name__)
 
-import re
-
 from tinyagentos.cluster.capabilities import hardware_to_targets, potential_capabilities as _potential_capabilities
 from tinyagentos.cluster.optimiser import ClusterOptimiser
 from tinyagentos.cluster.worker_auth import _HMACError, require_worker_hmac
@@ -900,7 +900,6 @@ async def deploy_backend(request: Request, name: str, body: DeployRequest):
             status_code=400,
         )
 
-    import httpx
     try:
         async with httpx.AsyncClient(timeout=620) as client:
             resp = await client.post(
@@ -935,7 +934,6 @@ async def worker_remote_command(request: Request, name: str, body: WorkerRemoteR
             status_code=403,
         )
 
-    import httpx
     try:
         async with httpx.AsyncClient(timeout=body.timeout + 5) as client:
             resp = await client.post(
@@ -1153,7 +1151,12 @@ async def update_worker(request: Request, name: str):
     worker = cluster.get_worker(name)
     if not worker:
         return JSONResponse({"error": f"Worker '{name}' not found"}, status_code=404)
-    if worker.status not in ("online", "draining"):
+    if worker.status == "draining":
+        return JSONResponse(
+            {"error": f"Worker '{name}' is already updating (status=draining)"},
+            status_code=409,
+        )
+    if worker.status != "online":
         return JSONResponse(
             {"error": f"Worker '{name}' is not online (status={worker.status})"},
             status_code=400,
@@ -1164,8 +1167,10 @@ async def update_worker(request: Request, name: str):
     if "error" in drain_result:
         return JSONResponse(drain_result, status_code=404)
 
-    # Step 2: Trigger the update-worker deploy command on the worker
-    import httpx
+    # Step 2: Trigger the update-worker deploy command on the worker.
+    # The deploy script runs the update and then restarts the worker
+    # service.  The HTTP connection will break when the worker restarts
+    # — that is expected and must not cancel the drain.
     deploy_result = None
     try:
         async with httpx.AsyncClient(timeout=620) as client:
@@ -1173,10 +1178,31 @@ async def update_worker(request: Request, name: str):
                 f"{worker.url}/api/worker/deploy",
                 json={"command": "update-worker"},
             )
+            resp.raise_for_status()
             deploy_result = resp.json()
+    except httpx.HTTPStatusError:
+        # The deploy endpoint returned a non-2xx status — real failure.
+        cluster.cancel_drain(name)
+        logger.exception("Worker '%s' update deploy returned HTTP error", name)
+        return JSONResponse(
+            {"error": f"Worker update deploy rejected by worker", "drain_cancelled": True},
+            status_code=502,
+        )
+    except (httpx.ConnectError, httpx.RemoteProtocolError, httpx.ReadError):
+        # Worker restarted mid-request — the connection dropped because
+        # the worker service stopped.  This is the normal update path.
+        # Keep the drain active; the worker will re-register and the
+        # monitor loop will auto-complete the drain when all leases
+        # are released (taOS #890).
+        logger.info(
+            "Worker '%s' disconnected during update — restart expected, "
+            "drain stays active until re-registration", name,
+        )
+        deploy_result = {"status": "acknowledged", "note": "worker disconnected — restart expected"}
     except Exception as exc:
-        # Deploy failed — cancel drain so worker can still serve traffic
+        # Unexpected failure — cancel drain so worker can still serve traffic.
         cluster.cancel_drain(name)
+        logger.exception("Worker '%s' update deploy failed unexpectedly", name)
         return JSONResponse(
             {"error": f"Worker update deploy failed: {exc}", "drain_cancelled": True},
             status_code=502,

From ad9d0738b66e06b9a6a983f034708f309fc14ed1 Mon Sep 17 00:00:00 2001
From: Hogne <227774406+hognek@users.noreply.github.com>
Date: Sat, 11 Jul 2026 00:52:08 +0200
Subject: [PATCH 4/4] fix(cluster): address re-review blockers on worker
 auto-update PR #1690
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

CRITICAL: Add _require_admin auth gate to /drain, /cancel-drain, /update
endpoints — all three were missing the admin-session check that every
other mutating cluster route uses.

WARNING — disconnect classification: Remove ConnectError from the
expected-restart tuple (connection-refused means unreachable, not
restarting) and add ReadTimeout (the real mid-restart signal).
ConnectError now cancels the drain and returns 502.

WARNING — lease-lock: Make drain_worker async and wrap the
force-release lease pops in async with self._lease_lock so
concurrent claim/release operations see a consistent view.

WARNING — task retention: Capture notification create_task results
in drain_worker and cancel_drain, retain in self._background_tasks
with add_done_callback(discard).

Minor — cancel_drain returns 409 (not 404) on state conflict.
Minor — bare f-string fix in update_worker error message.

Tests: 22/22 drain+update tests pass (9 unit + 13 HTTP), including
6 new tests (3 auth-gate, ConnectError-as-failure, ReadTimeout-as-success,
cancel_drain 409).
---
 tests/test_cluster.py          |  14 ++--
 tests/test_routes_cluster.py   | 116 +++++++++++++++++++++++++++++++++
 tinyagentos/cluster/manager.py |  31 +++++----
 tinyagentos/routes/cluster.py  |  29 +++++++--
 4 files changed, 166 insertions(+), 24 deletions(-)

diff --git a/tests/test_cluster.py b/tests/test_cluster.py
index 111b843b5..9e543b366 100644
--- a/tests/test_cluster.py
+++ b/tests/test_cluster.py
@@ -270,7 +270,7 @@ async def test_drain_worker_graceful_sets_draining_status(self):
         mgr = ClusterManager()
         await mgr.register_worker(_make_worker("gpu-box", capabilities=["chat"]))
 
-        result = mgr.drain_worker("gpu-box", graceful=True)
+        result = await mgr.drain_worker("gpu-box", graceful=True)
         assert result["worker"] == "gpu-box"
         assert result["previous_status"] == "online"
         assert result["status"] == "draining"
@@ -294,7 +294,7 @@ async def test_drain_worker_force_releases_leases_and_marks_offline(self):
         )
         mgr._leases["l_test1"] = lease
 
-        result = mgr.drain_worker("gpu-box", graceful=False)
+        result = await mgr.drain_worker("gpu-box", graceful=False)
         assert result["released_leases"] == 1
         assert result["status"] == "offline"
         assert mgr.get_worker("gpu-box").status == "offline"
@@ -303,7 +303,7 @@ async def test_drain_worker_force_releases_leases_and_marks_offline(self):
     async def test_drain_worker_unknown_returns_error(self):
         """Draining a non-existent worker returns error dict."""
         mgr = ClusterManager()
-        result = mgr.drain_worker("nonexistent")
+        result = await mgr.drain_worker("nonexistent")
         assert "error" in result
         assert result["worker"] == "nonexistent"
 
@@ -311,7 +311,7 @@ async def test_cancel_drain_returns_worker_to_online(self):
         """Cancel drain returns a draining worker back to online."""
         mgr = ClusterManager()
         await mgr.register_worker(_make_worker("gpu-box", capabilities=["chat"]))
-        mgr.drain_worker("gpu-box", graceful=True)
+        await mgr.drain_worker("gpu-box", graceful=True)
         assert mgr.get_worker("gpu-box").status == "draining"
 
         result = mgr.cancel_drain("gpu-box")
@@ -339,7 +339,7 @@ async def test_draining_workers_excluded_from_routing(self):
         mgr = ClusterManager()
         await mgr.register_worker(_make_worker("online-gpu", capabilities=["chat"], load=0.1))
         await mgr.register_worker(_make_worker("draining-gpu", capabilities=["chat"], load=0.0))
-        mgr.drain_worker("draining-gpu", graceful=True)
+        await mgr.drain_worker("draining-gpu", graceful=True)
 
         result = mgr.get_workers_for_capability("chat")
         assert len(result) == 1
@@ -354,7 +354,7 @@ async def test_draining_workers_excluded_from_catalog(self):
         draining.backends = [{"name": "b2", "type": "ollama", "url": "u", "capabilities": ["chat"], "models": [{"name": "m2"}]}]
         await mgr.register_worker(online)
         await mgr.register_worker(draining)
-        mgr.drain_worker("draining", graceful=True)
+        await mgr.drain_worker("draining", graceful=True)
 
         out = mgr.aggregate_catalog()
         assert [w["name"] for w in out["workers"]] == ["online"]
@@ -369,7 +369,7 @@ async def test_draining_workers_excluded_from_lease_claim(self):
         worker = mgr.get_worker("draining-gpu")
         worker.free_vram_mb = 8000
 
-        mgr.drain_worker("draining-gpu", graceful=True)
+        await mgr.drain_worker("draining-gpu", graceful=True)
 
         lease = await mgr.claim_lease(
             resource_id="draining-gpu:gpu-cuda-0",
diff --git a/tests/test_routes_cluster.py b/tests/test_routes_cluster.py
index d804d2b43..7dd0df6a0 100644
--- a/tests/test_routes_cluster.py
+++ b/tests/test_routes_cluster.py
@@ -683,3 +683,119 @@ async def test_update_worker_deploy_unexpected_error(client, app):
     assert "deploy failed" in data["error"]
     assert data["drain_cancelled"] is True
     assert cluster.get_worker("gpu-box").status == "online"
+
+
+# ── update_worker auth gate tests ───────────────────────────────────────
+
+
+@pytest.mark.asyncio
+async def test_drain_requires_admin(app):
+    """403 when drain endpoint is called without an admin session."""
+    from httpx import ASGITransport
+
+    async with AsyncClient(
+        transport=ASGITransport(app=app), base_url="http://test"
+    ) as noauth_client:
+        resp = await noauth_client.post("/api/cluster/workers/gpu-box/drain", json={})
+    assert resp.status_code in (401, 403), f"got {resp.status_code}: {resp.text}"
+
+
+@pytest.mark.asyncio
+async def test_cancel_drain_requires_admin(app):
+    """403 when cancel-drain endpoint is called without an admin session."""
+    from httpx import ASGITransport
+
+    async with AsyncClient(
+        transport=ASGITransport(app=app), base_url="http://test"
+    ) as noauth_client:
+        resp = await noauth_client.post("/api/cluster/workers/gpu-box/cancel-drain")
+    assert resp.status_code in (401, 403), f"got {resp.status_code}: {resp.text}"
+
+
+@pytest.mark.asyncio
+async def test_update_worker_requires_admin(app):
+    """403 when update endpoint is called without an admin session."""
+    from httpx import ASGITransport
+
+    async with AsyncClient(
+        transport=ASGITransport(app=app), base_url="http://test"
+    ) as noauth_client:
+        resp = await noauth_client.post("/api/cluster/workers/gpu-box/update")
+    assert resp.status_code in (401, 403), f"got {resp.status_code}: {resp.text}"
+
+
+# ── update_worker disconnect classification tests ────────────────────────
+
+
+@pytest.mark.asyncio
+async def test_update_worker_deploy_connect_error(client, app):
+    """502 when deploy endpoint raises ConnectError — connection refused means
+    the worker is genuinely unreachable, not restarting.  Drain is cancelled."""
+    from tinyagentos.cluster.worker_protocol import WorkerInfo
+    import httpx as _httpx
+    cluster = app.state.cluster_manager
+    cluster._workers["gpu-box"] = WorkerInfo(  # noqa: SLF001
+        name="gpu-box", url="http://gpu:9000", status="online",
+    )
+
+    mock_client = AsyncMock()
+    mock_client.__aenter__.return_value = mock_client
+    mock_client.__aexit__.return_value = None
+    mock_client.post.side_effect = _httpx.ConnectError(
+        "Connection refused"
+    )
+
+    with patch("tinyagentos.routes.cluster.httpx.AsyncClient", return_value=mock_client):
+        resp = await client.post("/api/cluster/workers/gpu-box/update")
+
+    assert resp.status_code == 502
+    data = resp.json()
+    assert "connection refused" in data["error"].lower()
+    assert data["drain_cancelled"] is True
+    # Drain must be cancelled — worker back online
+    assert cluster.get_worker("gpu-box").status == "online"
+
+
+@pytest.mark.asyncio
+async def test_update_worker_deploy_read_timeout(client, app):
+    """200 when deploy endpoint raises ReadTimeout — worker restarted
+    mid-request, which is the normal update path.  Drain stays active."""
+    from tinyagentos.cluster.worker_protocol import WorkerInfo
+    import httpx as _httpx
+    cluster = app.state.cluster_manager
+    cluster._workers["gpu-box"] = WorkerInfo(  # noqa: SLF001
+        name="gpu-box", url="http://gpu:9000", status="online",
+    )
+
+    mock_client = AsyncMock()
+    mock_client.__aenter__.return_value = mock_client
+    mock_client.__aexit__.return_value = None
+    mock_client.post.side_effect = _httpx.ReadTimeout(
+        "Read timed out", request=AsyncMock()
+    )
+
+    with patch("tinyagentos.routes.cluster.httpx.AsyncClient", return_value=mock_client):
+        resp = await client.post("/api/cluster/workers/gpu-box/update")
+
+    assert resp.status_code == 200
+    data = resp.json()
+    assert data["worker"] == "gpu-box"
+    assert data["status"] == "updating"
+    assert data["deploy"]["status"] == "acknowledged"
+    assert "restart expected" in data["deploy"]["note"]
+    # Drain must NOT be cancelled — worker should still be draining
+    assert cluster.get_worker("gpu-box").status == "draining"
+
+
+@pytest.mark.asyncio
+async def test_cancel_drain_state_conflict_returns_409(client, app):
+    """409 when cancel-drain is called on a worker not in draining state."""
+    from tinyagentos.cluster.worker_protocol import WorkerInfo
+    cluster = app.state.cluster_manager
+    cluster._workers["gpu-box"] = WorkerInfo(  # noqa: SLF001
+        name="gpu-box", url="http://gpu:9000", status="online",
+    )
+
+    resp = await client.post("/api/cluster/workers/gpu-box/cancel-drain")
+    assert resp.status_code == 409
+    assert "not draining" in resp.json()["error"]
diff --git a/tinyagentos/cluster/manager.py b/tinyagentos/cluster/manager.py
index 9a7ee8b08..d3a3a12b3 100644
--- a/tinyagentos/cluster/manager.py
+++ b/tinyagentos/cluster/manager.py
@@ -397,7 +397,7 @@ def _sweep_expired_leases(self):
 
     # ── taOS #890: graceful worker drain for auto-update ───────────────
 
-    def drain_worker(self, name: str, graceful: bool = True) -> dict:
+    async def drain_worker(self, name: str, graceful: bool = True) -> dict:
         """Begin draining a worker — gracefully detach without dropping tasks.
 
         When ``graceful=True`` (the default):
@@ -425,15 +425,18 @@ def drain_worker(self, name: str, graceful: bool = True) -> dict:
         released = 0
 
         if not graceful:
-            # Force-release all leases for this worker
-            lids = [
-                lid for lid, lease in self._leases.items()
-                if (parsed := self._parse_resource_id(lease.resource_id))
-                and parsed[0] == name
-            ]
-            for lid in lids:
-                self._leases.pop(lid, None)
-                released += 1
+            # Force-release all leases for this worker.
+            # Lease pops must be serialised under _lease_lock so concurrent
+            # claim/release operations see a consistent view (taOS #1690).
+            async with self._lease_lock:
+                lids = [
+                    lid for lid, lease in self._leases.items()
+                    if (parsed := self._parse_resource_id(lease.resource_id))
+                    and parsed[0] == name
+                ]
+                for lid in lids:
+                    self._leases.pop(lid, None)
+                    released += 1
             logger.info("Worker '%s' drain: force-released %d leases", name, released)
             worker.status = "offline"
 
@@ -443,7 +446,7 @@ def drain_worker(self, name: str, graceful: bool = True) -> dict:
                 f"{'Tasks will complete before detach.' if graceful else 'All leases released immediately.'}"
             )
             try:
-                asyncio.get_running_loop().create_task(
+                task = asyncio.create_task(
                     self._notifications.emit_event(
                         "worker.drain",
                         f"Worker '{name}' draining",
@@ -451,6 +454,8 @@ def drain_worker(self, name: str, graceful: bool = True) -> dict:
                         level="info",
                     )
                 )
+                self._background_tasks.add(task)
+                task.add_done_callback(self._background_tasks.discard)
             except RuntimeError:
                 pass
 
@@ -477,7 +482,7 @@ def cancel_drain(self, name: str) -> dict:
 
         if self._notifications:
             try:
-                asyncio.get_running_loop().create_task(
+                task = asyncio.create_task(
                     self._notifications.emit_event(
                         "worker.online",
                         f"Worker '{name}' drain cancelled",
@@ -485,6 +490,8 @@ def cancel_drain(self, name: str) -> dict:
                         level="info",
                     )
                 )
+                self._background_tasks.add(task)
+                task.add_done_callback(self._background_tasks.discard)
             except RuntimeError:
                 pass
 
diff --git a/tinyagentos/routes/cluster.py b/tinyagentos/routes/cluster.py
index 8db2df984..eff507200 100644
--- a/tinyagentos/routes/cluster.py
+++ b/tinyagentos/routes/cluster.py
@@ -1119,8 +1119,11 @@ async def drain_worker(request: Request, name: str, body: DrainRequest = DrainRe
     When ``graceful=false``, all leases are force-released and the worker
     is marked offline immediately.
     """
+    ok, err = _require_admin(request)
+    if not ok:
+        return err
     cluster = request.app.state.cluster_manager
-    result = cluster.drain_worker(name, graceful=body.graceful)
+    result = await cluster.drain_worker(name, graceful=body.graceful)
     if "error" in result:
         return JSONResponse(result, status_code=404)
     return result
@@ -1129,10 +1132,13 @@ async def drain_worker(request: Request, name: str, body: DrainRequest = DrainRe
 @router.post("/api/cluster/workers/{name}/cancel-drain")
 async def cancel_drain(request: Request, name: str):
     """Cancel an in-progress drain and return the worker to online status."""
+    ok, err = _require_admin(request)
+    if not ok:
+        return err
     cluster = request.app.state.cluster_manager
     result = cluster.cancel_drain(name)
     if "error" in result:
-        return JSONResponse(result, status_code=404)
+        return JSONResponse(result, status_code=409)
     return result
 
 
@@ -1147,6 +1153,9 @@ async def update_worker(request: Request, name: str):
     4. Worker restarts, re-registers, returns to "online".
     5. Routing resumes automatically once the worker is back online.
     """
+    ok, err = _require_admin(request)
+    if not ok:
+        return err
     cluster = request.app.state.cluster_manager
     worker = cluster.get_worker(name)
     if not worker:
@@ -1163,7 +1172,7 @@ async def update_worker(request: Request, name: str):
         )
 
     # Step 1: Begin draining
-    drain_result = cluster.drain_worker(name, graceful=True)
+    drain_result = await cluster.drain_worker(name, graceful=True)
     if "error" in drain_result:
         return JSONResponse(drain_result, status_code=404)
 
@@ -1185,10 +1194,10 @@ async def update_worker(request: Request, name: str):
         cluster.cancel_drain(name)
         logger.exception("Worker '%s' update deploy returned HTTP error", name)
         return JSONResponse(
-            {"error": f"Worker update deploy rejected by worker", "drain_cancelled": True},
+            {"error": "Worker update deploy rejected by worker", "drain_cancelled": True},
             status_code=502,
         )
-    except (httpx.ConnectError, httpx.RemoteProtocolError, httpx.ReadError):
+    except (httpx.RemoteProtocolError, httpx.ReadError, httpx.ReadTimeout):
         # Worker restarted mid-request — the connection dropped because
         # the worker service stopped.  This is the normal update path.
         # Keep the drain active; the worker will re-register and the
@@ -1199,6 +1208,16 @@ async def update_worker(request: Request, name: str):
             "drain stays active until re-registration", name,
         )
         deploy_result = {"status": "acknowledged", "note": "worker disconnected — restart expected"}
+    except httpx.ConnectError:
+        # Connection-refused / DNS failure — the worker is genuinely
+        # unreachable, not restarting.  Cancel the drain so the worker
+        # can be retried.
+        cluster.cancel_drain(name)
+        logger.exception("Worker '%s' update deploy failed — connection refused", name)
+        return JSONResponse(
+            {"error": "Worker unreachable — connection refused", "drain_cancelled": True},
+            status_code=502,
+        )
     except Exception as exc:
         # Unexpected failure — cancel drain so worker can still serve traffic.
         cluster.cancel_drain(name)