Expose worker idle time so callers can shut down inactive workers — Closes #253

wool/proto/wire.proto

@@ -86,6 +86,7 @@ message Nack {
 service Worker {
     rpc dispatch (stream Request) returns (stream Response);
     rpc stop (StopRequest) returns (Void);
+    rpc idle (Void) returns (IdleTime);
 }
 
 message Request {
@@ -117,6 +118,13 @@ message StopRequest {
     float timeout = 1;
 }
 
+message IdleTime {
+    // Seconds the worker has been continuously idle: time since the
+    // in-flight task set last emptied (startup counts as empty).
+    // 0 while any task is in flight; measured on a monotonic clock.
+    double seconds = 1;
+}
+
 message WorkerMetadata {
     // Unique identifier for the worker instance
     string uid = 1;

wool/src/wool/__init__.py

@@ -48,6 +48,7 @@
 from wool.runtime.worker.base import Worker
 from wool.runtime.worker.base import WorkerFactory
 from wool.runtime.worker.base import WorkerLike
+from wool.runtime.worker.connection import IdleUnavailable
 from wool.runtime.worker.connection import RpcError
 from wool.runtime.worker.connection import TransientRpcError
 from wool.runtime.worker.connection import UnexpectedResponse
@@ -110,6 +111,7 @@
     "DiscoveryPublisherLike",
     "DiscoverySubscriberLike",
     "Factory",
+    "IdleUnavailable",
     "IneffectiveLeaseWarning",
     "IneffectiveQuorumTimeoutWarning",
     "LanDiscovery",

wool/src/wool/protocol/__init__.py

@@ -13,6 +13,7 @@
 from wool.protocol._wire import ChainManifest as ChainManifest
 from wool.protocol._wire import ChannelOptions as ChannelOptions
 from wool.protocol._wire import ContextVar as ContextVar
+from wool.protocol._wire import IdleTime as IdleTime
 from wool.protocol._wire import Message as Message
 from wool.protocol._wire import Nack as Nack
 from wool.protocol._wire import Request as Request
@@ -35,6 +36,7 @@
     "ChainManifest",
     "ChannelOptions",
     "ContextVar",
+    "IdleTime",
     "Message",
     "Nack",
     "Request",

wool/src/wool/protocol/_wire.py

@@ -14,6 +14,7 @@
     from wool.protocol.wire_pb2 import ChainManifest
     from wool.protocol.wire_pb2 import ChannelOptions
     from wool.protocol.wire_pb2 import ContextVar
+    from wool.protocol.wire_pb2 import IdleTime
     from wool.protocol.wire_pb2 import Message
     from wool.protocol.wire_pb2 import Nack
     from wool.protocol.wire_pb2 import Request
@@ -57,6 +58,7 @@ def __call__(servicer, server) -> None: ...
     "ChainManifest",
     "ChannelOptions",
     "ContextVar",
+    "IdleTime",
     "Message",
     "Nack",
     "Request",

wool/src/wool/runtime/worker/README.md

@@ -139,14 +139,14 @@ class ResilientWorker(LocalWorker):
         await super()._start(timeout=timeout)
         self._monitor_task = asyncio.create_task(self._monitor())
 
-    async def _stop(self, timeout=None):
+    async def _stop(self, grace=None):
         if self._monitor_task:
             self._monitor_task.cancel()
             try:
                 await self._monitor_task
             except asyncio.CancelledError:
                 pass
-        await super()._stop(timeout=timeout)
+        await super()._stop(grace=grace)
 
     async def _restart(self):
         """Replace the dead process, reusing the original port."""

wool/src/wool/runtime/worker/base.py

@@ -3,6 +3,7 @@
 import functools
 import inspect
 import uuid
+import warnings
 from abc import ABC
 from abc import abstractmethod
 from dataclasses import dataclass
@@ -295,11 +296,15 @@ async def start(self, *, timeout: float | None = None):
         """
         ...
 
-    async def stop(self, *, timeout: float | None = None):
+    async def stop(self, *, grace: float | None = None, timeout: float | None = None):
         """Stop the worker and unregister it from the pool.
 
+        :param grace:
+            The worker's shutdown grace period in seconds — how long to
+            wait for in-flight tasks to drain before cancelling them.
         :param timeout:
-            Maximum time in seconds to wait for worker shutdown.
+            Deprecated alias for ``grace``, retained for backwards
+            compatibility; passing it emits a ``DeprecationWarning``.
         :raises RuntimeError:
             If the worker has not been started.
         """
@@ -326,7 +331,7 @@ async def _start(self, timeout):
                 # Start your worker process
                 self._info = WorkerMetadata(...)
 
-            async def _stop(self, timeout):
+            async def _stop(self, grace):
                 # Clean shutdown
                 ...
 
@@ -406,17 +411,33 @@ async def start(self, *, timeout: float | None = None):
         assert self._info
 
     @final
-    async def stop(self, *, timeout: float | None = None):
+    async def stop(self, *, grace: float | None = None, timeout: float | None = None):
         """Stop the worker and unregister it from the pool.
 
         This method is a final implementation that calls the abstract
         `_stop` method to gracefully shut down the worker process and
         unregister it from the registrar service.
+
+        :param grace:
+            The worker's shutdown grace period in seconds — how long to
+            wait for in-flight tasks to drain before cancelling them.
+        :param timeout:
+            Deprecated alias for ``grace``, retained for backwards
+            compatibility; passing it emits a ``DeprecationWarning``.
         """
+        if timeout is not None:
+            warnings.warn(
+                "The 'timeout' parameter of Worker.stop is deprecated; "
+                "use 'grace' instead.",
+                DeprecationWarning,
+                stacklevel=2,
+            )
+            if grace is None:
+                grace = timeout
         if not self._started:
             raise RuntimeError("Worker has not been started")
         try:
-            await self._stop(timeout)
+            await self._stop(grace)
         finally:
             self._started = False
 
@@ -433,7 +454,7 @@ async def _start(self, timeout: float | None):
         ...
 
     @abstractmethod
-    async def _stop(self, timeout: float | None):
+    async def _stop(self, grace: float | None):
         """Implementation-specific worker shutdown logic.
 
         Subclasses must implement this method to handle the graceful

wool/src/wool/runtime/worker/connection.py

@@ -17,6 +17,7 @@
 
 import wool
 from wool import protocol
+from wool.exceptions import WoolError
 from wool.runtime.resourcepool import ResourcePool
 from wool.runtime.routine.task import Task
 from wool.runtime.serializer import Serializer
@@ -121,9 +122,29 @@ class TransientRpcError(RpcError):
     """
 
 
+# public
+class IdleUnavailable(WoolError):
+    """Raised when a worker does not implement the idle RPC.
+
+    A worker that predates the idle capability answers the idle RPC
+    with gRPC ``UNIMPLEMENTED``; `WorkerConnection.idle_time`
+    translates that into this typed signal so a polling client can
+    detect an old worker and treat idle reporting as unavailable,
+    distinct from a transient hiccup or an unhealthy peer. It descends
+    from `wool.WoolError` — not `RpcError` — because an absent
+    capability is not an RPC-health fault, so ``except RpcError`` does
+    not catch it.
+    """
+
+
 # public
 class WorkerConnection:
-    """gRPC connection to a worker for task dispatch.
+    """Direct single-worker control surface over a pooled gRPC channel.
+
+    Exposes `dispatch` (task execution), `idle_time` (poll the worker's
+    continuous idle duration), and `stop` (shut the remote worker
+    down). ``close`` is distinct: it releases this connection's local
+    pooled channel, whereas `stop` terminates the remote worker.
 
     Acquires pooled gRPC channels keyed by ``(target, credentials,
     options)``.  Each `dispatch` call obtains a reference-counted
@@ -329,6 +350,76 @@ async def close(self):
             except KeyError:
                 pass
 
+    async def idle_time(self, *, timeout: float | None = None) -> float:
+        """Query how long the remote worker has been continuously idle.
+
+        Returns the worker's reported continuous idle duration; see
+        `WorkerService.idle` for how idle is defined and measured. The
+        call draws a channel from this connection's pool, inheriting
+        its credential and secure-channel handling.
+
+        :param timeout:
+            gRPC call deadline in seconds for the poll. ``None``
+            applies no deadline.
+        :returns:
+            The worker's continuous idle duration in seconds.
+        :raises ValueError:
+            If ``timeout`` is not positive.
+        :raises IdleUnavailable:
+            If the worker predates the idle RPC (gRPC ``UNIMPLEMENTED``).
+        :raises TransientRpcError:
+            If the worker returns a transient RPC error
+            (``UNAVAILABLE``, ``DEADLINE_EXCEEDED``, or
+            ``RESOURCE_EXHAUSTED``).
+        :raises RpcError:
+            If the worker returns any other non-transient RPC error.
+        """
+        if timeout is not None and timeout <= 0:
+            raise ValueError("Idle timeout must be positive")
+        async with _channel_pool.get(self._key) as channel:
+            try:
+                response = await channel.stub.idle(protocol.Void(), timeout=timeout)
+            except grpc.RpcError as error:
+                code = error.code()
+                details = error.details() or str(error)
+                if code == grpc.StatusCode.UNIMPLEMENTED:
+                    raise IdleUnavailable(
+                        "Worker does not implement idle reporting"
+                    ) from error
+                if code in self.TRANSIENT_ERRORS:
+                    raise TransientRpcError(code, details) from error
+                raise RpcError(code, details) from error
+        return response.seconds
+
+    async def stop(self, *, grace: float | None = None) -> None:
+        """Stop the remote worker.
+
+        Sends the stop RPC over a channel drawn from this connection's
+        pool, inheriting its credential and secure-channel handling.
+
+        :param grace:
+            The worker's shutdown grace period in seconds, carried in
+            the stop request. ``None`` (the wire default of ``0``)
+            cancels in-flight tasks immediately; a positive value
+            waits that long for a graceful drain; a negative value
+            waits indefinitely.
+        :raises TransientRpcError:
+            If the worker returns a transient RPC error
+            (``UNAVAILABLE``, ``DEADLINE_EXCEEDED``, or
+            ``RESOURCE_EXHAUSTED``).
+        :raises RpcError:
+            If the worker returns any other non-transient RPC error.
+        """
+        async with _channel_pool.get(self._key) as channel:
+            try:
+                await channel.stub.stop(protocol.StopRequest(timeout=grace))
+            except grpc.RpcError as error:
+                code = error.code()
+                details = error.details() or str(error)
+                if code in self.TRANSIENT_ERRORS:
+                    raise TransientRpcError(code, details) from error
+                raise RpcError(code, details) from error
+
     async def _handshake(
         self,
         call: _DispatchCall,

wool/src/wool/runtime/worker/local.py

@@ -4,12 +4,10 @@
 from typing import TYPE_CHECKING
 from typing import Any
 
-import grpc.aio
-
-from wool import protocol
 from wool.runtime.worker.auth import WorkerCredentials
 from wool.runtime.worker.base import Worker
 from wool.runtime.worker.base import WorkerOptions
+from wool.runtime.worker.connection import WorkerConnection
 from wool.runtime.worker.process import WorkerProcess
 
 if TYPE_CHECKING:
@@ -134,30 +132,29 @@ async def _start(self, timeout: float | None):
         if self._info is None:
             raise RuntimeError("Worker process failed to start - no metadata")
 
-    async def _stop(self, timeout: float | None):
+    async def _stop(self, grace: float | None):
         """Stop the worker process and unregister it from the pool.
 
         Unregisters the worker from the registrar service, gracefully
         shuts down the worker process using a gRPC stop request, and cleans
         up the registrar service. If graceful shutdown fails, the process
         is forcefully terminated.
 
-        For workers configured with :class:`WorkerCredentials`, uses
-        the client credentials to establish a secure connection for the
-        stop operation. For insecure workers, uses an insecure channel.
+        Credential and secure-channel handling for the stop RPC lives
+        on `WorkerConnection.stop`.
         """
         if self._worker_process.is_alive():
             assert self.address
 
-            # Create appropriate channel based on available credentials
-            if self._credentials is not None:
-                credentials = self._credentials.client_credentials()
-                channel = grpc.aio.secure_channel(self.address, credentials)
-            else:
-                channel = grpc.aio.insecure_channel(self.address)
-
+            # Route the stop RPC through WorkerConnection; ``close``
+            # releases the pooled channel this stop acquired.
+            credentials = (
+                self._credentials.client_credentials()
+                if self._credentials is not None
+                else None
+            )
+            connection = WorkerConnection(self.address, credentials=credentials)
             try:
-                stub = protocol.WorkerStub(channel)
-                await stub.stop(protocol.StopRequest(timeout=timeout))
+                await connection.stop(grace=grace)
             finally:
-                await channel.close()
+                await connection.close()

wool/src/wool/runtime/worker/pool.py

@@ -648,7 +648,7 @@ async def start(worker):
 
                 async def stop(worker):
                     await publisher_svc.publish("worker-dropped", worker.metadata)
-                    await worker.stop(timeout=self._shutdown_timeout)
+                    await worker.stop(grace=self._shutdown_timeout)
 
                 task = asyncio.create_task(start(worker))
                 tasks.append(task)