[cross-repo from server#137] Worker Heartbeats / Status surface: every SDK emits periodic heartbeat; CLI+UI list workers per task queue (Temporal-parity)

Author: durable-workflow-ops

src/durable_workflow/client.py

@@ -609,12 +609,17 @@ class WorkerDescription:
     last_heartbeat_at: str | None = None
     registered_at: str | None = None
     updated_at: str | None = None
+    task_slots: dict[str, int | None] | None = None
+    process_metrics: dict[str, Any] | None = None
+    heartbeat_interval_seconds: int | None = None
     raw: dict[str, Any] | None = None
 
     @classmethod
     def from_dict(cls, data: dict[str, Any], *, worker_id: str | None = None) -> WorkerDescription:
         workflow_types = data.get("supported_workflow_types")
         activity_types = data.get("supported_activity_types")
+        task_slots = data.get("task_slots")
+        process_metrics = data.get("process_metrics")
 
         return cls(
             worker_id=data.get("worker_id", worker_id or ""),
@@ -631,6 +636,13 @@ def from_dict(cls, data: dict[str, Any], *, worker_id: str | None = None) -> Wor
             last_heartbeat_at=data.get("last_heartbeat_at"),
             registered_at=data.get("registered_at"),
             updated_at=data.get("updated_at"),
+            task_slots=task_slots if isinstance(task_slots, dict) else None,
+            process_metrics=process_metrics if isinstance(process_metrics, dict) else None,
+            heartbeat_interval_seconds=(
+                int(data["heartbeat_interval_seconds"])
+                if isinstance(data.get("heartbeat_interval_seconds"), int)
+                else None
+            ),
             raw=data,
         )
 
@@ -641,6 +653,7 @@ class WorkerList:
 
     namespace: str | None
     workers: list[WorkerDescription]
+    stale_after_seconds: int | None = None
 
 
 @dataclass
@@ -1818,6 +1831,11 @@ async def list_workers(
                 for item in items
                 if isinstance(item, dict)
             ],
+            stale_after_seconds=(
+                int(data["stale_after_seconds"])
+                if isinstance(data.get("stale_after_seconds"), int)
+                else None
+            ),
         )
 
     async def describe_worker(self, worker_id: str) -> WorkerDescription:
@@ -2613,6 +2631,50 @@ async def register_worker(
             body["max_concurrent_activity_tasks"] = max_concurrent_activity_tasks
         return await self._request("POST", "/worker/register", worker=True, json=body)
 
+    async def heartbeat_worker(
+        self,
+        *,
+        worker_id: str,
+        task_slots: dict[str, int] | None = None,
+        process_metrics: dict[str, Any] | None = None,
+        heartbeat_interval_seconds: int | None = None,
+    ) -> Any:
+        """Send a worker-fleet heartbeat to refresh liveness and report state.
+
+        Workers should call this on a steady cadence (default 60s, advertised
+        by the server in the register/heartbeat acknowledgement) so operators
+        can answer "what workers are polling task queue X right now, what's
+        their slot capacity, when did each last check in" via the worker
+        management API, the CLI worker listing, and the operator Worker
+        Status view.
+
+        ``task_slots`` is an optional dict with any subset of
+        ``workflow_available``, ``activity_available``, ``session_available``
+        — the count of currently free slots for each family. The server
+        clamps each value into ``[0, max_concurrent_*]``.
+
+        ``process_metrics`` is an optional dict with any subset of
+        ``cpu_percent``, ``memory_bytes``, ``process_uptime_seconds``,
+        ``process_id``, and ``host`` — the SDK reports only what it has
+        cheap access to, and the server records exactly what was reported.
+
+        Returns the server acknowledgement, which includes the advertised
+        ``heartbeat_interval_seconds`` and ``stale_after_seconds`` so the
+        worker can adapt its cadence on the fly.
+
+        Most applications create a :class:`~durable_workflow.Worker`, which
+        runs this on a background asyncio task — call this directly only when
+        driving the worker protocol by hand (smoke tests, custom runtimes).
+        """
+        body: dict[str, Any] = {"worker_id": worker_id}
+        if task_slots:
+            body["task_slots"] = task_slots
+        if process_metrics:
+            body["process_metrics"] = process_metrics
+        if heartbeat_interval_seconds is not None:
+            body["heartbeat_interval_seconds"] = heartbeat_interval_seconds
+        return await self._request("POST", "/worker/heartbeat", worker=True, json=body)
+
     async def poll_workflow_task(
         self, *, worker_id: str, task_queue: str, timeout: float = 35.0
     ) -> Any:

src/durable_workflow/worker.py

@@ -20,6 +20,7 @@
 import hashlib
 import inspect
 import logging
+import sys
 import time
 import traceback
 import uuid
@@ -290,6 +291,7 @@ def __init__(
         max_concurrent_workflow_tasks: int = 10,
         max_concurrent_activity_tasks: int = 10,
         shutdown_timeout: float = 30.0,
+        heartbeat_interval: float = 60.0,
         metrics: MetricsRecorder | None = None,
         interceptors: Iterable[WorkerInterceptor] = (),
     ) -> None:
@@ -313,6 +315,8 @@ def __init__(
             raise ValueError("max_concurrent_workflow_tasks must be at least 1")
         if max_concurrent_activity_tasks < 1:
             raise ValueError("max_concurrent_activity_tasks must be at least 1")
+        if heartbeat_interval <= 0:
+            raise ValueError("heartbeat_interval must be positive")
 
         self._poll_timeout = poll_timeout
         self.max_concurrent_workflow_tasks = max_concurrent_workflow_tasks
@@ -323,6 +327,13 @@ def __init__(
         self._shutdown_timeout = shutdown_timeout
         self._in_flight: set[asyncio.Task[Any]] = set()
         self._query_tasks_supported = False
+        # In-flight slot accounting feeds the periodic heartbeat so operators
+        # see free-slot counts without the worker having to re-derive them at
+        # shutdown. Counters are bumped/decremented around dispatch.
+        self._workflow_inflight = 0
+        self._activity_inflight = 0
+        self._heartbeat_interval = float(heartbeat_interval)
+        self._process_started_at = time.time()
         configured_metrics = metrics if metrics is not None else getattr(client, "metrics", NOOP_METRICS)
         self.metrics: MetricsRecorder = configured_metrics or NOOP_METRICS
         self.interceptors = tuple(interceptors)
@@ -388,7 +399,7 @@ async def _register(self) -> None:
             _manifest_version(info.get("worker_protocol")),
         )
 
-        await self.client.register_worker(
+        ack = await self.client.register_worker(
             worker_id=self.worker_id,
             task_queue=self.task_queue,
             supported_workflow_types=list(self.workflows),
@@ -398,6 +409,14 @@ async def _register(self) -> None:
             max_concurrent_activity_tasks=self.max_concurrent_activity_tasks,
             build_id=self.build_id,
         )
+        # Adapt to the server-advertised cadence when present so a cluster
+        # can pin the worker fleet's heartbeat beat without each worker
+        # passing the cadence explicitly. Falls back to the constructor
+        # value when the server has not advertised a cadence.
+        if isinstance(ack, dict):
+            advertised = ack.get("heartbeat_interval_seconds")
+            if isinstance(advertised, int) and advertised > 0:
+                self._heartbeat_interval = float(advertised)
         log.info("worker %s registered on %s", self.worker_id, self.task_queue)
 
     async def _run_workflow_task(self, task: dict[str, Any]) -> list[dict[str, Any]] | None:
@@ -993,12 +1012,14 @@ async def _poll_workflow_tasks(self) -> None:
     async def _dispatch_workflow_task(self, task: dict[str, Any]) -> None:
         task_start = time.perf_counter()
         outcome = "error"
+        self._workflow_inflight += 1
         try:
             commands = await self._run_workflow_task(task)
             outcome = "completed" if commands is not None else "failed"
         except Exception:
             log.exception("unhandled error in workflow task execution")
         finally:
+            self._workflow_inflight = max(0, self._workflow_inflight - 1)
             self._record_task_metrics("workflow", outcome, time.perf_counter() - task_start)
             self._wf_semaphore.release()
 
@@ -1032,11 +1053,13 @@ async def _poll_activity_tasks(self) -> None:
     async def _dispatch_activity_task(self, task: dict[str, Any]) -> None:
         task_start = time.perf_counter()
         outcome = "error"
+        self._activity_inflight += 1
         try:
             outcome = await self._run_activity_task(task)
         except Exception:
             log.exception("unhandled error in activity task execution")
         finally:
+            self._activity_inflight = max(0, self._activity_inflight - 1)
             self._record_task_metrics("activity", outcome, time.perf_counter() - task_start)
             self._act_semaphore.release()
 
@@ -1076,12 +1099,96 @@ async def run(self) -> None:
         await self._register()
         wf_loop = asyncio.create_task(self._poll_workflow_tasks())
         act_loop = asyncio.create_task(self._poll_activity_tasks())
-        loops = [wf_loop, act_loop]
+        hb_loop = asyncio.create_task(self._heartbeat_loop())
+        loops = [wf_loop, act_loop, hb_loop]
         if self._query_tasks_supported:
             loops.append(asyncio.create_task(self._poll_query_tasks()))
         with contextlib.suppress(asyncio.CancelledError):
             await asyncio.gather(*loops)
 
+    async def _heartbeat_loop(self) -> None:
+        """Periodically refresh the server-side worker registration.
+
+        Reports current task-slot availability and basic process-level
+        metrics so the worker management API, CLI worker listing, and
+        Waterline Worker Status view can show free-slot counts and
+        process health alongside ``last_heartbeat_at``. Cadence is the
+        server-advertised ``heartbeat_interval_seconds`` (default 60s,
+        bounded to [1s, 1h] cluster-wide) so workers stop being
+        considered for task dispatch when they miss enough heartbeats.
+        """
+        while not self._stop.is_set():
+            try:
+                await asyncio.wait_for(self._stop.wait(), timeout=self._heartbeat_interval)
+            except asyncio.TimeoutError:
+                pass
+            if self._stop.is_set():
+                return
+            try:
+                ack = await self.client.heartbeat_worker(
+                    worker_id=self.worker_id,
+                    task_slots=self._current_task_slots(),
+                    process_metrics=self._current_process_metrics(),
+                )
+            except Exception as e:
+                log.warning("worker heartbeat failed: %s", e)
+                continue
+            if isinstance(ack, dict):
+                advertised = ack.get("heartbeat_interval_seconds")
+                if isinstance(advertised, int) and advertised > 0:
+                    self._heartbeat_interval = float(advertised)
+
+    def _current_task_slots(self) -> dict[str, int]:
+        return {
+            "workflow_available": max(
+                0, self.max_concurrent_workflow_tasks - self._workflow_inflight
+            ),
+            "activity_available": max(
+                0, self.max_concurrent_activity_tasks - self._activity_inflight
+            ),
+        }
+
+    def _current_process_metrics(self) -> dict[str, Any]:
+        import os
+        import socket
+
+        metrics: dict[str, Any] = {
+            "process_uptime_seconds": int(time.time() - self._process_started_at),
+            "process_id": os.getpid(),
+        }
+
+        try:
+            import resource
+
+            usage = resource.getrusage(resource.RUSAGE_SELF)
+            # ru_maxrss is kilobytes on Linux and bytes on macOS — normalize
+            # to bytes. The server stores whatever is sent so the units stay
+            # consistent across SDKs.
+            if sys.platform == "darwin":
+                metrics["memory_bytes"] = int(usage.ru_maxrss)
+            else:
+                metrics["memory_bytes"] = int(usage.ru_maxrss) * 1024
+
+            cpu_seconds = float(usage.ru_utime) + float(usage.ru_stime)
+            wall_seconds = max(0.001, time.time() - self._process_started_at)
+            metrics["cpu_percent"] = max(
+                0.0, min(100.0, round((cpu_seconds / wall_seconds) * 100.0, 2))
+            )
+        except (ImportError, OSError):
+            # `resource` is POSIX-only — Windows skips getrusage but still
+            # reports pid + uptime + host so the operator surface remains
+            # populated.
+            pass
+
+        try:
+            host = socket.gethostname()
+        except Exception:
+            host = ""
+        if isinstance(host, str) and host != "":
+            metrics["host"] = host[:255]
+
+        return metrics
+
     async def run_until(
         self,
         *,

tests/test_worker.py

@@ -90,6 +90,9 @@ async def echo_async_activity(val: str) -> str:
 def mock_client() -> AsyncMock:
     client = AsyncMock(spec=Client)
     client.register_worker = AsyncMock(return_value={"worker_id": "w1", "registered": True})
+    client.heartbeat_worker = AsyncMock(
+        return_value={"worker_id": "w1", "acknowledged": True, "heartbeat_interval_seconds": 60}
+    )
     client.poll_workflow_task = AsyncMock(return_value=None)
     client.poll_activity_task = AsyncMock(return_value=None)
     client.poll_query_task = AsyncMock(return_value=None)
@@ -1357,6 +1360,83 @@ async def test_run_skips_query_loop_without_query_task_capability(self, mock_cli
         mock_client.poll_query_task.assert_not_called()
 
 
+class TestWorkerHeartbeats:
+    @pytest.mark.asyncio
+    async def test_run_drives_periodic_heartbeats_with_slot_state(
+        self, mock_client: AsyncMock
+    ) -> None:
+        worker = Worker(
+            mock_client,
+            task_queue="q1",
+            workflows=[TestWorkflow],
+            activities=[echo_activity],
+            max_concurrent_workflow_tasks=4,
+            max_concurrent_activity_tasks=2,
+            poll_timeout=0.01,
+            heartbeat_interval=0.05,
+        )
+        run_task = asyncio.create_task(worker.run())
+        await asyncio.sleep(0.2)
+        await worker.stop()
+        run_task.cancel()
+        with contextlib.suppress(asyncio.CancelledError):
+            await run_task
+
+        assert mock_client.heartbeat_worker.call_count >= 1
+        kwargs = mock_client.heartbeat_worker.call_args.kwargs
+        assert kwargs["worker_id"] == worker.worker_id
+        assert kwargs["task_slots"]["workflow_available"] == 4
+        assert kwargs["task_slots"]["activity_available"] == 2
+        process_metrics = kwargs["process_metrics"]
+        assert "process_id" in process_metrics
+        assert process_metrics["process_id"] > 0
+        assert "process_uptime_seconds" in process_metrics
+
+    @pytest.mark.asyncio
+    async def test_register_adopts_server_advertised_heartbeat_cadence(
+        self, mock_client: AsyncMock
+    ) -> None:
+        mock_client.register_worker = AsyncMock(
+            return_value={
+                "worker_id": "w1",
+                "registered": True,
+                "heartbeat_interval_seconds": 7,
+            }
+        )
+        worker = Worker(
+            mock_client,
+            task_queue="q1",
+            workflows=[TestWorkflow],
+            activities=[echo_activity],
+            heartbeat_interval=120.0,
+        )
+        await worker._register()
+        assert worker._heartbeat_interval == 7.0
+
+    @pytest.mark.asyncio
+    async def test_heartbeat_loop_survives_transient_errors(
+        self, mock_client: AsyncMock
+    ) -> None:
+        mock_client.heartbeat_worker = AsyncMock(
+            side_effect=[RuntimeError("temporary"), {"acknowledged": True}]
+        )
+        worker = Worker(
+            mock_client,
+            task_queue="q1",
+            workflows=[TestWorkflow],
+            activities=[echo_activity],
+            poll_timeout=0.01,
+            heartbeat_interval=0.02,
+        )
+        run_task = asyncio.create_task(worker.run())
+        await asyncio.sleep(0.15)
+        await worker.stop()
+        run_task.cancel()
+        with contextlib.suppress(asyncio.CancelledError):
+            await run_task
+        assert mock_client.heartbeat_worker.call_count >= 2
+
+
 class TestRunUntil:
     @pytest.mark.asyncio
     async def test_run_until_returns_terminal_description(self, mock_client: AsyncMock) -> None: