ringo-task-queue 0.1.0.dev8__py3-none-macosx_13_0_arm64.whl

ringo_task_queue/worker.py

@@ -0,0 +1,1868 @@
+"""Asyncio worker implementation for the credit-controlled Work stream.
+
+The worker deliberately owns one reader and one writer for every Work call.
+Handlers only receive :class:`TaskContext`; lease tokens stay in this module.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import inspect
+import math
+import sys
+from collections import deque
+from dataclasses import dataclass, field
+from datetime import datetime, timedelta, timezone
+from typing import TYPE_CHECKING, Any, Awaitable, Callable, Iterable
+
+import grpc
+
+from . import _pb
+from ._proto.ringo.v1 import queue_pb2
+from ._supervisor_policy import TaskTimeouts
+from .concurrency import AdaptiveConcurrency, ResourceConcurrency, _Member
+from .errors import (
+    IncompatibleVersionError,
+    LeaseLostError,
+    NotFoundError,
+    PermanentTaskError,
+    RetryTaskError,
+    RingoError,
+    UnavailableError,
+    error_from_detail,
+    map_rpc_error,
+)
+from .models import Task
+
+if TYPE_CHECKING:
+    from ._supervisor_watchdog import _WorkerWatchdog
+
+
+def _current_runtime() -> Any:
+    # A normal SDK install does not import the process-supervision runtime.
+    module = sys.modules.get("ringo_task_queue._supervisor_runtime")
+    return None if module is None else module.current_runtime()
+
+
+_MAX_ERROR_BYTES = 256 * 1024
+
+#: Fixed on_idle monitor cadence (design §10.6). These are deliberately not
+#: exposed as Worker parameters: the poll interval and the continuous-idle
+#: quiet period are protocol-level constants of the idle-observation feature.
+_IDLE_POLL_SECONDS = 0.25
+_IDLE_QUIET_SECONDS = 1.0
+#: GetQueueActivity was introduced in protocol minor 1.
+_IDLE_MIN_PROTOCOL_MINOR = 1
+_ORDERED_REQUEST_IDS_MIN_PROTOCOL_MINOR = 6
+_MAX_REQUEST_ID = 2**64 - 1
+
+
+class _WatchdogFenced(Exception):
+    """A stopped deadline invocation abandoned during graceful drain."""
+
+
+def _consume_task_exception(task: asyncio.Future[Any]) -> None:
+    if not task.cancelled():
+        task.exception()
+
+
+def _safe_error(value: object) -> str:
+    """Return a valid UTF-8 error string no larger than 256 KiB."""
+    try:
+        text = str(value)
+    except Exception:
+        text = type(value).__name__
+    encoded = text.encode("utf-8", errors="replace")
+    if len(encoded) <= _MAX_ERROR_BYTES:
+        return encoded.decode("utf-8", errors="replace")
+    # Decode with replacement after truncating bytes; this is UTF-8 safe and
+    # never emits a broken surrogate or an over-sized protobuf string.
+    return encoded[:_MAX_ERROR_BYTES].decode("utf-8", errors="ignore")
+
+
+class TaskContext:
+    """Context supplied to a worker handler.
+
+    ``cancelled`` becomes true when the local lease is lost, the stream drops,
+    or the worker is closed.  ``report_progress`` is correlated on the Work
+    stream and never exposes the lease token to application code.
+    """
+
+    __slots__ = (
+        "_worker",
+        "_task",
+        "_attempt",
+        "_max_attempts",
+        "_worker_id",
+        "_assignment",
+        "_cancel_event",
+    )
+
+    def __init__(
+        self,
+        worker: "Worker",
+        task: Task,
+        attempt: int,
+        max_attempts: int,
+        worker_id: str,
+        assignment: "_Assignment",
+    ) -> None:
+        self._worker = worker
+        self._task = task
+        self._attempt = attempt
+        self._max_attempts = max_attempts
+        self._worker_id = worker_id
+        self._assignment = assignment
+        self._cancel_event = asyncio.Event()
+
+    @property
+    def task(self) -> Task:
+        return self._task
+
+    @property
+    def attempt(self) -> int:
+        return self._attempt
+
+    @property
+    def max_attempts(self) -> int:
+        return self._max_attempts
+
+    @property
+    def worker_id(self) -> str:
+        return self._worker_id
+
+    @property
+    def cancelled(self) -> bool:
+        return self._cancel_event.is_set()
+
+    @property
+    def is_cancelled(self) -> bool:
+        """Alias useful to code shared with SDKs using a verb-like name."""
+        return self.cancelled
+
+    async def wait_cancelled(self) -> None:
+        await self._cancel_event.wait()
+
+    def mark_progress(self) -> None:
+        """Mark local business progress without sending any RPC."""
+        watchdog = self._worker._watchdog
+        if watchdog is not None:
+            watchdog.mark_progress(self._assignment)
+
+    async def report_progress(
+        self, current: float, total: float, message: str = ""
+    ) -> None:
+        """Mark local progress, report it and await its correlated result."""
+        if self.cancelled:
+            raise LeaseLostError("task lease is no longer valid")
+        self.mark_progress()
+        await self._worker._report_progress(self._assignment, current, total, message)
+
+    def _cancel(self) -> None:
+        self._cancel_event.set()
+
+
+@dataclass
+class _Assignment:
+    task: Task
+    attempt: int
+    token: str
+    until: datetime
+    context: TaskContext | None = None
+    handler_task: asyncio.Task[Any] | None = None
+    renew_task: asyncio.Task[Any] | None = None
+    alive: bool = True
+    completion_started: bool = False
+    completion_success: bool = False
+    completion_congestion: bool = False
+    controller_released: bool = False
+    controller_release_task: asyncio.Task[None] | None = None
+    fixed_released: bool = False
+    started: bool = False
+    session_epoch: int = 0
+    session: asyncio.Event | None = None
+    renew_deadline: float | None = None
+    watchdog_expired: bool = False
+
+    @property
+    def identity(self) -> tuple[str, int, str]:
+        return (self.task.id, self.attempt, self.token)
+
+    @property
+    def ref(self) -> queue_pb2.LeaseRef:
+        return queue_pb2.LeaseRef(
+            task_id=self.task.id,
+            attempt=self.attempt,
+            lease_token=self.token,
+        )
+
+
+@dataclass
+class _WorkSource:
+    """One Work call owns its correlation and transport lifetime.
+
+    Reader/writer tasks keep this object even after a replacement becomes the
+    current source. No receipt, EOF or timeout from an old call can touch the
+    replacement's pending Futures, registration or transport.
+    """
+
+    epoch: int
+    call: Any
+    minor: int = 0
+    token: str = ""
+    ordered_request_ids: bool = False
+    done: asyncio.Event = field(default_factory=asyncio.Event)
+    registered: asyncio.Event = field(default_factory=asyncio.Event)
+    outbox: asyncio.Queue[queue_pb2.WorkRequest | None] = field(
+        default_factory=asyncio.Queue
+    )
+    pending: dict[int, asyncio.Future[queue_pb2.WorkResult]] = field(
+        default_factory=dict
+    )
+    next_request_id: int = 1
+    error: BaseException | None = None
+    reader: asyncio.Task[Any] | None = None
+    writer: asyncio.Task[Any] | None = None
+    publisher: asyncio.Task[None] | None = None
+    cleanup: asyncio.Task[None] | None = None
+    # Diagnostics only: accepted credit never participates in admission itself.
+    accepted_credit: int = 0
+    credit_generation: int = -1
+
+
+class Worker:
+    """A credit-controlled asynchronous task worker.
+
+    ``run`` is intentionally a blocking coroutine: it remains active until
+    ``close`` is called.  The client registers workers and drains them before
+    closing its channel.  When ``on_idle`` is configured, ``run`` additionally
+    observes queue activity and executes the callback inline at most once.
+    """
+
+    def __init__(
+        self,
+        queue: Any,
+        handler: Callable[[Task, TaskContext], Awaitable[Any]],
+        *,
+        task_types: Iterable[str],
+        concurrency: int | ResourceConcurrency | AdaptiveConcurrency,
+        lease_duration: timedelta,
+        grace_period: timedelta,
+        on_idle: Callable[[], Any] | None = None,
+        weight: int = 1,
+        task_timeouts: TaskTimeouts | None = None,
+    ) -> None:
+        if isinstance(concurrency, bool) or not isinstance(
+            concurrency, (int, ResourceConcurrency, AdaptiveConcurrency)
+        ):
+            raise ValueError(
+                "concurrency must be an integer, ResourceConcurrency, or AdaptiveConcurrency"
+            )
+        if isinstance(concurrency, int) and concurrency < 1:
+            raise ValueError("initial fixed concurrency must be >= 1")
+        if isinstance(weight, bool) or not isinstance(weight, int) or weight < 1:
+            raise ValueError("weight must be a positive integer")
+        self._queue = queue
+        self._client = queue._client
+        self._handler = handler
+        self._task_types = list(task_types)
+        self._controller: ResourceConcurrency | AdaptiveConcurrency | None = (
+            concurrency if not isinstance(concurrency, int) else None
+        )
+        self._member: _Member | None = None
+        self._concurrency = (
+            concurrency
+            if isinstance(concurrency, int)
+            else max(1, concurrency.capacity)
+        )
+        self._weight = weight
+        self._lease_duration = lease_duration
+        self._grace_period = max(0.0, grace_period.total_seconds())
+        self._worker_id = self._new_worker_id()
+        self._on_idle = on_idle
+        # on_idle monitor state (design §10.6). ``_idle_candidate`` arms only
+        # after WorkerRegistered plus the initial credit of the *current*
+        # session; every reset condition disarms it or clears the zero timer.
+        self._idle_candidate = False
+        self._idle_zero_since: float | None = None
+        self._idle_fired = False
+        #: Session token captured at WorkerRegistered + initial credit. The
+        #: candidate is only observable while the connection still carries
+        #: this token: a unary-triggered renegotiation rotates the token
+        #: without touching the Work stream and must disarm the monitor.
+        self._idle_session_token: str | None = None
+        #: Set when the on_idle callback raised: run() must re-raise verbatim
+        #: instead of counting it as a Work-stream reconnect failure.
+        self._idle_error: BaseException | None = None
+
+        # close_started means no new stream may be opened, while force_cancel
+        # is raised only after the graceful drain window expires.  Existing
+        # handlers and renewals remain fully live between those two points.
+        self._close_started = asyncio.Event()
+        self._force_cancel = asyncio.Event()
+        self._run_started = False
+        self._run_finished = asyncio.Event()
+        self._close_finished = asyncio.Event()
+        self._command_timeout = 30.0
+        self._close_lock = asyncio.Lock()
+        self._close_task: asyncio.Task[None] | None = None
+        self._source: _WorkSource | None = None
+        # Recovery history is distinct from a live source's error. A clean EOF
+        # remains clean, while exhausted recovery can report its earlier cause.
+        self._last_stream_error: BaseException | None = None
+        # A successor lease can arrive before the predecessor's NACK receipt
+        # is consumed. Only an identical complete LeaseRef is a duplicate.
+        self._inflight: dict[tuple[str, int, str], _Assignment] = {}
+        self._waiting_assignments: deque[_Assignment] = deque()
+        self._capacity_tasks: set[asyncio.Task[Any]] = set()
+        # Handler lifetime, not just lease-map membership, owns local capacity.
+        # Old-stream handlers that swallow cancellation remain here until their
+        # coroutine actually settles.
+        self._handler_tasks: set[asyncio.Task[Any]] = set()
+        self._fixed_running = 0
+        # Every capacity update shares one publisher mutex and a monotonically
+        # increasing session epoch.  A late task from an old stream can never
+        # publish credit onto a replacement stream or after close.
+        self._publisher_lock = asyncio.Lock()
+        self._session_epoch = 0
+        self._publisher_closed = False
+        self._fixed_capacity_generation = 0
+        self._fixed_publication_task: asyncio.Task[None] | None = None
+        self._fixed_publication_source: tuple[int, asyncio.Event] | None = None
+        self._fixed_publication_waiters: list[tuple[int, asyncio.Future[None]]] = []
+        self._fixed_dynamic_target = False
+        if task_timeouts is not None and not isinstance(task_timeouts, TaskTimeouts):
+            raise ValueError("task_timeouts must be TaskTimeouts or None")
+        self._runtime = _current_runtime()
+        self._task_timeouts = (
+            self._runtime.task_timeouts
+            if task_timeouts is None and self._runtime is not None
+            else task_timeouts
+        )
+        self._watchdog: _WorkerWatchdog | None = None
+        self._watchdog_error: RuntimeError | None = None
+        if (
+            self._runtime is not None
+            and self._runtime.policy.no_start_timeout is not None
+        ) or (self._task_timeouts is not None and self._task_timeouts.enabled):
+            from ._supervisor_watchdog import _WorkerWatchdog
+
+            self._watchdog = _WorkerWatchdog(
+                self, self._task_timeouts or TaskTimeouts(), self._runtime
+            )
+        if self._controller is not None:
+            self._member = self._controller._attach(self, weight)
+
+    @staticmethod
+    def _new_worker_id() -> str:
+        # Importing here avoids coupling worker construction to client internals.
+        from .client import new_worker_id
+
+        return new_worker_id()
+
+    # Read-only private views keep diagnostics and existing benchmark probes
+    # useful without maintaining a second set of mutable transport state.
+    @property
+    def _session_done(self) -> asyncio.Event | None:
+        return self._source.done if self._source is not None else None
+
+    @property
+    def _call(self) -> Any:
+        return self._source.call if self._source is not None else None
+
+    @property
+    def _outbox(self) -> asyncio.Queue[queue_pb2.WorkRequest | None] | None:
+        return self._source.outbox if self._source is not None else None
+
+    @property
+    def _pending(self) -> dict[int, asyncio.Future[queue_pb2.WorkResult]]:
+        return self._source.pending if self._source is not None else {}
+
+    @property
+    def _registered(self) -> asyncio.Event | None:
+        return self._source.registered if self._source is not None else None
+
+    @property
+    def _reader_task(self) -> asyncio.Task[Any] | None:
+        return self._source.reader if self._source is not None else None
+
+    @property
+    def _writer_task(self) -> asyncio.Task[Any] | None:
+        return self._source.writer if self._source is not None else None
+
+    @property
+    def _stream_error(self) -> BaseException | None:
+        return (
+            self._source.error if self._source is not None else self._last_stream_error
+        )
+
+    @property
+    def worker_id(self) -> str:
+        return self._worker_id
+
+    @property
+    def concurrency(self) -> int:
+        """Current fixed target, or the controller's current capacity."""
+        if self._controller is not None:
+            return self._controller.capacity
+        return self._concurrency
+
+    @property
+    def weight(self) -> int:
+        return self._weight
+
+    async def set_concurrency(self, concurrency: int) -> None:
+        """Resize a fixed worker without cancelling running handlers."""
+        if self._controller is not None:
+            raise ValueError(
+                "a Worker attached to a shared controller cannot resize itself"
+            )
+        if isinstance(concurrency, bool) or not isinstance(concurrency, int):
+            raise ValueError("concurrency must be an integer")
+        if concurrency < 0 or concurrency > 2**31 - 1:
+            raise ValueError("concurrency must be between 0 and 2**31-1")
+        if self._run_started:
+            conn = self._client._require_conn()
+            if conn._session_minor < 2:
+                raise IncompatibleVersionError(
+                    "runtime fixed concurrency requires protocol minor >= 2; "
+                    f"negotiated minor is {conn._session_minor}"
+                )
+        self._concurrency = concurrency
+        self._fixed_dynamic_target = True
+        self._start_waiting_assignments()
+        if self._session_done is not None and not self._session_done.is_set():
+            await self._publish_fixed_capacity()
+
+    async def run(self) -> None:
+        """Run until :meth:`close` is called, reconnecting boundedly."""
+        if self._run_started:
+            raise RuntimeError("worker.run() may only be called once")
+        self._run_started = True
+        failures = 0
+        try:
+            if self._runtime is not None:
+                self._runtime.register_worker(self)
+            if self._watchdog is not None:
+                self._watchdog.start()
+            conn = self._client._require_conn()
+            if (
+                self._controller is not None or self._fixed_dynamic_target
+            ) and conn._session_minor < 2:
+                raise IncompatibleVersionError(
+                    "dynamic concurrency requires protocol minor >= 2; "
+                    f"negotiated minor is {conn._session_minor}"
+                )
+            if self._on_idle is not None:
+                # GetQueueActivity arrived in protocol minor 1; fail fast
+                # instead of silently running without idle observation.
+                conn = self._client._require_conn()
+                if conn._session_minor < _IDLE_MIN_PROTOCOL_MINOR:
+                    raise IncompatibleVersionError(
+                        "on_idle requires a server speaking protocol minor"
+                        f" >= {_IDLE_MIN_PROTOCOL_MINOR}; negotiated minor"
+                        f" is {conn._session_minor}"
+                    )
+            while not self._close_started.is_set():
+                try:
+                    await self._run_session_monitored()
+                except asyncio.CancelledError:
+                    if not self._close_started.is_set():
+                        raise
+                    break
+                except Exception as exc:
+                    if self._idle_error is not None:
+                        # An on_idle callback failure is not a transport
+                        # error: it must not be retried, reconnected, or
+                        # converted into a task NACK.
+                        raise
+                    if isinstance(exc, IncompatibleVersionError):
+                        # A reconnect may have renegotiated a protocol
+                        # downgrade that makes on_idle unservable; surface
+                        # it instead of burning the reconnect budget.
+                        raise
+                    self._last_stream_error = exc
+                    failures += 1
+                    if self._close_started.is_set():
+                        break
+                if self._close_started.is_set():
+                    break
+                # A stream loss invalidates local completion attempts, but does
+                # not send terminal commands for leases that may still exist.
+                await self._cancel_assignments()
+                conn = self._client._require_conn()
+                before_token = conn._session_token
+                await conn._recover_best_effort()
+                if conn._session_token == before_token:
+                    failures += 1
+                else:
+                    failures = 0
+                if failures >= 3:
+                    raise self._stream_error or UnavailableError(
+                        "bounded Work stream recovery exhausted"
+                    )
+                if not self._close_started.is_set():
+                    await asyncio.sleep(0)
+            if self._watchdog_error is not None:
+                raise self._watchdog_error
+        finally:
+            try:
+                await self._cancel_assignments()
+            finally:
+                try:
+                    if self._watchdog is not None:
+                        await self._watchdog.stop()
+                finally:
+                    if self._runtime is not None:
+                        self._runtime.unregister_worker(self)
+                    self._run_finished.set()
+
+    async def close(self) -> None:
+        """Drain and stop once; caller cancellation cannot cancel cleanup."""
+        async with self._close_lock:
+            if self._close_task is None:
+                self._close_started.set()
+                # The controller must linearize detach and withdraw its probe
+                # before this gate closes publication; the close path uses the
+                # still-live stream to await that correlated WorkResult.
+                self._idle_candidate = False
+                self._idle_zero_since = None
+                self._idle_session_token = None
+                self._close_task = asyncio.create_task(
+                    self._close_impl(), name=f"ringo-close-{self._worker_id}"
+                )
+                self._close_task.add_done_callback(_consume_task_exception)
+            close_task = self._close_task
+        await asyncio.shield(close_task)
+
+    async def _close_impl(self) -> None:
+        try:
+            if self._controller is not None and self._member is not None:
+                await self._controller._detach(self._member)
+            self._publisher_closed = True
+            # Drain keeps this Work session alive for handler ACK/renewal.
+            # Close is not a session replacement: changing its epoch here
+            # would fence valid origin-bound terminal commands during grace.
+            session = self._session_done
+            # Drain is best effort. A disconnected stream cannot safely report
+            # any terminal result, and leases are deliberately left to expire.
+            if session is not None and not session.is_set():
+                try:
+                    await asyncio.wait_for(
+                        self._send_command(
+                            queue_pb2.WorkRequest(
+                                drain=queue_pb2.WorkerDrain(begin=True),
+                            ),
+                            expected_epoch=self._session_epoch,
+                            expected_session=session,
+                        ),
+                        timeout=max(
+                            0.1,
+                            min(self._command_timeout, self._grace_period or 1.0),
+                        ),
+                    )
+                except Exception:
+                    pass
+            if self._grace_period:
+                try:
+                    await asyncio.wait_for(
+                        self._wait_handlers(), timeout=self._grace_period
+                    )
+                except asyncio.TimeoutError:
+                    pass
+        finally:
+            # From this point handlers and renewals are cancelled; no terminal
+            # frame is generated by their cancellation paths.
+            self._force_cancel.set()
+            try:
+                await self._cancel_assignments()
+                await self._stop_session()
+                self._client._unregister_worker(self)
+                if self._run_started:
+                    await self._run_finished.wait()
+            finally:
+                self._close_finished.set()
+
+    async def _wait_handlers(self) -> None:
+        tasks = list(self._handler_tasks)
+        if tasks:
+            await asyncio.gather(*tasks, return_exceptions=True)
+        # Done callbacks release permits asynchronously. Drain those releases
+        # too, but never include buffered leases that have not entered a handler.
+        if self._capacity_tasks:
+            await asyncio.gather(
+                *(asyncio.shield(task) for task in list(self._capacity_tasks)),
+                return_exceptions=True,
+            )
+
+    async def _cancel_assignments(self) -> None:
+        # Local assignments were reset, so the idle observation restarts too.
+        self._idle_candidate = False
+        self._idle_zero_since = None
+        self._idle_session_token = None
+        assignments = list(self._inflight.values())
+        self._inflight.clear()
+        self._waiting_assignments.clear()
+        renewals: list[asyncio.Task[Any]] = []
+        for assignment in assignments:
+            self._invalidate_assignment(assignment)
+            handler = assignment.handler_task
+            if handler is not None and not handler.done():
+                # Cancellation is advisory.  Never let a handler that catches
+                # CancelledError block reconnect or forced close; its local
+                # capacity remains occupied until its coroutine really ends.
+                handler.cancel()
+            renewal = assignment.renew_task
+            if renewal is not None and not renewal.done():
+                renewal.cancel()
+                renewals.append(renewal)
+        if renewals:
+            await asyncio.gather(*renewals, return_exceptions=True)
+        for assignment in assignments:
+            await self._release_capacity_if_settled(assignment)
+
+    def _invalidate_assignment(self, assignment: _Assignment) -> None:
+        if not assignment.alive:
+            return
+        assignment.alive = False
+        if assignment.context is not None:
+            assignment.context._cancel()
+
+    async def _release_controller_permit(
+        self,
+        assignment: _Assignment,
+        *,
+        success: bool = False,
+        congestion: bool = False,
+    ) -> None:
+        controller, member = self._controller, self._member
+        if controller is None or member is None or assignment.controller_released:
+            return
+        if assignment.controller_release_task is None:
+
+            async def release() -> None:
+                await controller._release_running(
+                    member, success=success, congestion=congestion
+                )
+                # The shielded task owns completion, even if every waiter was
+                # cancelled while the pool's resize lock held the decrement.
+                assignment.controller_released = True
+
+            task = asyncio.create_task(release())
+            assignment.controller_release_task = task
+            self._capacity_tasks.add(task)
+            task.add_done_callback(self._capacity_tasks.discard)
+            task.add_done_callback(_consume_task_exception)
+        # Grace timeout/caller cancellation must not consume the once flag
+        # while the pool's resize lock still blocks the actual decrement.
+        await asyncio.shield(assignment.controller_release_task)
+
+    async def _release_capacity_if_settled(self, assignment: _Assignment) -> None:
+        """Release capacity only as the application handler settles."""
+        if assignment.alive:
+            return
+        handler = assignment.handler_task
+        if handler is not None and not handler.done():
+            return
+        if self._controller is not None:
+            await self._release_controller_permit(
+                assignment,
+                success=assignment.completion_success,
+                congestion=assignment.completion_congestion,
+            )
+            return
+        if assignment.fixed_released:
+            return
+        assignment.fixed_released = True
+        if assignment.started:
+            self._fixed_running = max(0, self._fixed_running - 1)
+        self._start_waiting_assignments()
+        if not self._close_started.is_set():
+            await self._publish_fixed_capacity()
+
+    async def _run_session_monitored(self) -> None:
+        """Run one Work session, observing queue activity when on_idle is set.
+
+        Without ``on_idle`` this is exactly ``_run_session()``. With it, the
+        poll loop is executed inline by the ``run()`` coroutine itself (never
+        in a child task) so a callback exception — including ``SystemExit`` —
+        unwinds verbatim through ``await worker.run()``.
+        """
+        if self._on_idle is None:
+            await self._run_session()
+            return
+        session_task = asyncio.create_task(
+            self._run_session(), name="ringo-worker-session"
+        )
+        try:
+            while True:
+                done, _pending = await asyncio.wait(
+                    (session_task,), timeout=_IDLE_POLL_SECONDS
+                )
+                if done:
+                    # Propagate the session outcome (result or stream error).
+                    return await session_task
+                if await self._idle_poll_once():
+                    # Single-shot fired and returned normally: park on the
+                    # session instead of waking every poll interval.
+                    return await session_task
+        finally:
+            if not session_task.done():
+                session_task.cancel()
+            await asyncio.gather(session_task, return_exceptions=True)
+
+    async def _idle_poll_once(self) -> bool:
+        """One read-only, retry-safe activity observation tick.
+
+        Returns True exactly when the callback fired and returned normally.
+        """
+        callback = self._on_idle
+        if (
+            callback is None
+            or self._idle_fired
+            or not self._idle_candidate
+            or self._close_started.is_set()
+            or self._force_cancel.is_set()
+        ):
+            return False
+        # Local assignments make the queue non-idle for this worker.
+        if self._inflight or self._handler_tasks:
+            self._idle_zero_since = None
+            return False
+        session = self._session_done
+        if session is None or session.is_set():
+            return False
+        # The observation is bound to the token captured at registration: a
+        # unary-triggered renegotiation rotates the token without restarting
+        # the Work stream. End this session so the run loop re-registers the
+        # stream under the fresh token instead of observing a stale session.
+        conn = self._client._require_conn()
+        if conn._session_token != self._idle_session_token:
+            self._end_stale_idle_session(session)
+            return False
+        request = queue_pb2.GetQueueActivityRequest(
+            queue=self._queue.name, task_types=self._task_types
+        )
+        try:
+            response = await conn.call("GetQueueActivity", request, retry_safe=True)
+        except asyncio.CancelledError:
+            raise
+        except NotFoundError:
+            # A never-enqueued queue has no durable rows. This is a conclusive
+            # zero for idle observation only; public queries keep NOT_FOUND.
+            response = queue_pb2.GetQueueActivityResponse(active_task_count=0)
+        except Exception:
+            # A failed observation restarts the quiet window; a later poll
+            # may succeed again without tearing the session down.
+            self._idle_zero_since = None
+            return False
+        # The session or its token may have changed while the query was in
+        # flight; a stale response must not advance or complete the window.
+        if (
+            not self._idle_candidate
+            or self._session_done is not session
+            or session.is_set()
+        ):
+            self._idle_zero_since = None
+            return False
+        if self._client._require_conn()._session_token != self._idle_session_token:
+            self._end_stale_idle_session(session)
+            return False
+        if response.active_task_count > 0:
+            self._idle_zero_since = None
+            return False
+        now = asyncio.get_running_loop().time()
+        if self._idle_zero_since is None:
+            self._idle_zero_since = now
+            return False
+        if now - self._idle_zero_since < _IDLE_QUIET_SECONDS:
+            return False
+        # Final atomic gate: no await separates these checks from marking
+        # fired and invoking the callback, so an assignment landing in a
+        # query in-flight window cannot produce a false trigger.
+        if (
+            self._inflight
+            or self._handler_tasks
+            or self._close_started.is_set()
+            or self._force_cancel.is_set()
+        ):
+            self._idle_zero_since = None
+            return False
+        # Single-shot: mark before invoking so the callback can never
+        # re-enter the monitor, and so a normal return never fires twice.
+        self._idle_fired = True
+        try:
+            result = callback()
+            if inspect.isawaitable(result):
+                await result
+        except BaseException as exc:
+            # Tag callback failures so run() propagates them verbatim
+            # instead of feeding them to the reconnect loop.
+            self._idle_error = exc
+            raise
+        return True
+
+    def _end_stale_idle_session(self, session: asyncio.Event) -> None:
+        """End the current Work session after a token rotation.
+
+        Setting ``session_done`` is idempotent and is observed as a clean
+        stream end: ``_run_session`` returns without a stream error, the run
+        loop reconnects, and the next session re-registers (and re-arms the
+        idle candidate) under the fresh token. It is neither a callback
+        failure nor a reconnect-worthy stream error.
+        """
+        source = self._source
+        if source is None or source.done is not session:
+            return
+        self._idle_candidate = False
+        self._idle_zero_since = None
+        self._end_source(source)
+
+    def _end_source(
+        self, source: _WorkSource, error: BaseException | None = None
+    ) -> None:
+        """Synchronously end exactly one source before any cleanup wait.
+
+        Failing pending Futures here replaces per-command session-wait tasks.
+        A graceful Worker.close merely drains; it calls this only when the
+        transport actually ends or the grace period has expired.
+        """
+        if error is not None and source.error is None:
+            source.error = error
+            if self._source is source:
+                self._last_stream_error = error
+        source.done.set()
+        failure = source.error or UnavailableError("Work stream closed")
+        for future in source.pending.values():
+            if not future.done():
+                future.set_exception(failure)
+        source.pending.clear()
+        try:
+            source.call.cancel()
+        except Exception:
+            pass
+
+    async def _run_session(self) -> None:
+        conn = self._client._require_conn()
+        self._session_epoch += 1
+        # A new session invalidates any in-flight idle observation.
+        self._idle_candidate = False
+        self._idle_zero_since = None
+        self._idle_session_token = None
+        if self._on_idle is not None and conn._session_minor < _IDLE_MIN_PROTOCOL_MINOR:
+            # A reconnect may have renegotiated a downgrade to minor 0.
+            # Fail fast before opening the Work stream; run() surfaces this
+            # verbatim instead of counting it as a reconnect failure.
+            raise IncompatibleVersionError(
+                "on_idle requires a server speaking protocol minor"
+                f" >= {_IDLE_MIN_PROTOCOL_MINOR}; negotiated minor"
+                f" is {conn._session_minor}"
+            )
+        reg = queue_pb2.WorkerRegistration(
+            queue=self._queue.name,
+            worker_id=self._worker_id,
+            task_types=self._task_types,
+            max_concurrency=max(1, self._concurrency),
+        )
+        if self._lease_duration > timedelta(0):
+            reg.lease_duration.CopyFrom(_pb.td_to_dur(self._lease_duration))
+        async with conn._lock:
+            stub = conn._require_stub()
+            reg.context.CopyFrom(conn.request_context())
+            # Mode belongs to this registration snapshot, not to a later unary
+            # negotiation that may rotate the connection's token and minor.
+            minor = reg.context.protocol_minor
+            if (
+                self._controller is not None or self._fixed_dynamic_target
+            ) and minor < 2:
+                raise IncompatibleVersionError(
+                    "dynamic concurrency requires protocol minor >= 2; "
+                    f"negotiated minor is {minor}"
+                )
+            if self._on_idle is not None and minor < _IDLE_MIN_PROTOCOL_MINOR:
+                raise IncompatibleVersionError(
+                    "on_idle requires a server speaking protocol minor"
+                    f" >= {_IDLE_MIN_PROTOCOL_MINOR}; negotiated minor is {minor}"
+                )
+            ordered = minor >= _ORDERED_REQUEST_IDS_MIN_PROTOCOL_MINOR
+            if ordered:
+                reg.ordered_request_ids = True
+            call = stub.Work(metadata=conn._metadata())
+        source = _WorkSource(
+            self._session_epoch,
+            call,
+            minor=minor,
+            token=reg.context.session_token,
+            ordered_request_ids=ordered,
+        )
+        self._source = source
+        self._last_stream_error = None
+        source.writer = asyncio.create_task(
+            self._write_loop(source), name="ringo-worker-writer"
+        )
+        source.reader = asyncio.create_task(
+            self._read_loop(source), name="ringo-worker-reader"
+        )
+        try:
+            self._publish_frame(source, queue_pb2.WorkRequest(register=reg))
+            await self._wait_event_or_session(source.registered, source)
+            if self._force_cancel.is_set():
+                return
+            # Initial credit is sent only after WorkerRegistered was observed.
+            if self._controller is None:
+                await self._publish_fixed_capacity()
+            else:
+                self._controller._check_loop()
+                await self._controller._start_member(self._member)
+            if self._on_idle is not None and not self._close_started.is_set():
+                # Eligible only from the first/current WorkerRegistered plus
+                # a successful initial credit of this session. The candidate
+                # is bound to the token this stream was registered with: a
+                # later unary renegotiation that rotates the token disarms
+                # the monitor until the next session.
+                self._idle_session_token = source.token
+                self._idle_candidate = True
+                self._idle_zero_since = None
+            await source.done.wait()
+            if source.error is not None and not self._close_started.is_set():
+                raise source.error
+        finally:
+            await self._stop_session(source)
+            if self._source is source:
+                self._source = None
+
+    async def _wait_event_or_session(
+        self, event: asyncio.Event, source: _WorkSource
+    ) -> None:
+        # Registration is a cold-path Event race. Commands await only their
+        # private Future; every source-end path fails those Futures directly.
+        event_task = asyncio.create_task(event.wait())
+        done_task = asyncio.create_task(source.done.wait())
+        try:
+            done, _ = await asyncio.wait(
+                (event_task, done_task), return_when=asyncio.FIRST_COMPLETED
+            )
+            if done_task in done and not event.is_set():
+                raise source.error or UnavailableError("Work stream closed")
+        finally:
+            for task in (event_task, done_task):
+                if not task.done():
+                    task.cancel()
+            await asyncio.gather(event_task, done_task, return_exceptions=True)
+
+    async def _stop_session(self, source: _WorkSource | None = None) -> None:
+        source = self._source if source is None else source
+        if source is None:
+            return
+        # Fail command Futures before any lock/join can block teardown. One
+        # shielded owner survives repeated caller cancellation and concurrent
+        # close/run cleanup, including cancellation before its first step.
+        self._end_source(source)
+        if source.cleanup is None:
+            task = asyncio.create_task(
+                self._stop_source(source),
+                name=f"ringo-stop-{self._worker_id}-{source.epoch}",
+            )
+            source.cleanup = task
+            self._capacity_tasks.add(task)
+            task.add_done_callback(self._capacity_tasks.discard)
+            task.add_done_callback(_consume_task_exception)
+        await asyncio.shield(source.cleanup)
+
+    async def _stop_source(self, source: _WorkSource) -> None:
+        source.outbox.put_nowait(None)
+        tasks = [
+            task
+            for task in (source.reader, source.writer, source.publisher)
+            if task is not None and task is not asyncio.current_task()
+        ]
+        for task in tasks:
+            task.cancel()
+        controller, member = self._controller, self._member
+        if controller is not None and member is not None and self._source is source:
+            # The run loop joins this cleanup before admitting its replacement.
+            # Only unused old-stream credit is reclaimed; running handlers keep
+            # their permits until their real coroutine settles.
+            await controller._stream_lost(member)
+        if tasks:
+            await asyncio.gather(*tasks, return_exceptions=True)
+
+    def _publish_frame(self, source: _WorkSource, frame: queue_pb2.WorkRequest) -> None:
+        """Synchronous FIFO insertion; callers own validation and correlation."""
+        source.outbox.put_nowait(frame)
+
+    async def _write_loop(self, source: _WorkSource) -> None:
+        try:
+            while not source.done.is_set():
+                frame = await source.outbox.get()
+                if frame is None or source.done.is_set():
+                    try:
+                        done_writing = getattr(source.call, "done_writing", None)
+                        if done_writing is not None:
+                            await done_writing()
+                    except Exception:
+                        pass
+                    return
+                await source.call.write(frame)
+        except asyncio.CancelledError:
+            raise
+        except Exception as exc:
+            self._end_source(source, self._map_transport(exc))
+        finally:
+            self._end_source(source)
+
+    async def _read_loop(self, source: _WorkSource) -> None:
+        try:
+            async for response in source.call:
+                if source.done.is_set():
+                    return
+                body = response.WhichOneof("body")
+                if body == "registered":
+                    if self._source is not source:
+                        continue
+                    registered = response.registered
+                    if registered.ordered_request_ids != source.ordered_request_ids:
+                        raise IncompatibleVersionError(
+                            "Work registration did not confirm the requested ID mode"
+                        )
+                    if registered.HasField("lease_duration"):
+                        self._lease_duration = (
+                            _pb.dur_to_td(registered.lease_duration)
+                            or self._lease_duration
+                        )
+                    if (
+                        self._controller is None
+                        and not self._fixed_dynamic_target
+                        and registered.max_concurrency
+                    ):
+                        self._concurrency = registered.max_concurrency
+                    source.registered.set()
+                elif body == "task":
+                    await self._on_assignment(response.task, source=source)
+                elif body == "result":
+                    result = response.result
+                    future = source.pending.pop(result.request_id, None)
+                    if future is not None and not future.done():
+                        future.set_result(result)
+                # Unknown/late results and heartbeats keep their existing
+                # non-fatal behavior; they never correlate into another source.
+        except asyncio.CancelledError:
+            raise
+        except Exception as exc:
+            self._end_source(source, self._map_transport(exc))
+        finally:
+            self._end_source(source)
+
+    @staticmethod
+    def _map_transport(exc: BaseException) -> BaseException:
+        if isinstance(exc, RingoError):
+            return exc
+        if isinstance(exc, grpc.aio.AioRpcError):
+            return map_rpc_error(exc)
+        return UnavailableError(f"Work stream failed: {_safe_error(exc)}")
+
+    async def _on_assignment(
+        self, leased: queue_pb2.LeasedTask, *, source: _WorkSource | None = None
+    ) -> None:
+        source = self._source if source is None else source
+        if (
+            source is None
+            or self._source is not source
+            or source.done.is_set()
+            or self._force_cancel.is_set()
+            or not leased.HasField("task")
+            or not leased.HasField("lease")
+        ):
+            return
+        task = Task.from_proto(leased.task)
+        until = _pb.ts_to_dt(leased.lease.until)
+        assignment = _Assignment(
+            task=task,
+            attempt=leased.lease.attempt,
+            token=leased.lease.lease_token,
+            until=until,
+            renew_deadline=asyncio.get_running_loop().time()
+            + (until - datetime.now(timezone.utc)).total_seconds(),
+        )
+        context = TaskContext(
+            self,
+            task,
+            assignment.attempt,
+            task.max_attempts,
+            self._worker_id,
+            assignment,
+        )
+        assignment.context = context
+        # The predecessor may still be waiting for its successful NACK result
+        # when a successor arrives. Do not confuse that valid new lease with
+        # a duplicate, and never overwrite either lease's cleanup ownership.
+        if assignment.identity in self._inflight:
+            return
+        assignment.session_epoch = source.epoch
+        assignment.session = source.done
+        if self._controller is not None and self._member is not None:
+            # Consuming the reserved probe is the linearization point at which
+            # an assignment becomes a running controller permit.
+            if not await self._controller._probe_consumed(self._member):
+                # A withdrawn probe must not become a new handler.  The lease
+                # remains server-owned and will naturally expire if this was a
+                # transport race.
+                return
+            if self._source is not source or source.done.is_set():
+                assignment.alive = False
+                context._cancel()
+                await self._release_controller_permit(assignment)
+                return
+        # A complete short task can arrive and ACK between two activity polls.
+        # Acceptance itself restarts quiet time; old/duplicate/rejected frames
+        # returned above must never reset the current source's observation.
+        self._idle_zero_since = None
+        self._inflight[assignment.identity] = assignment
+        assignment.renew_task = asyncio.create_task(
+            self._renew_loop(assignment),
+            name=f"ringo-renew-{task.id}-{assignment.attempt}",
+        )
+        if self._controller is None:
+            # Absolute credit updates can race assignments already in transit.
+            # Retain and renew such leases, but gate actual handler entry locally.
+            self._waiting_assignments.append(assignment)
+            self._fixed_capacity_generation += 1
+            self._start_waiting_assignments()
+        else:
+            self._start_handler(assignment)
+
+    def _start_waiting_assignments(self) -> None:
+        if (
+            self._controller is not None
+            or self._close_started.is_set()
+            or self._force_cancel.is_set()
+            or self._session_done is None
+            or self._session_done.is_set()
+        ):
+            return
+        while self._waiting_assignments and self._fixed_running < self._concurrency:
+            assignment = self._waiting_assignments.popleft()
+            if not assignment.alive or assignment.session_epoch != self._session_epoch:
+                continue
+            expired = (
+                assignment.renew_deadline <= asyncio.get_running_loop().time()
+                if assignment.renew_deadline is not None
+                else assignment.until <= datetime.now(timezone.utc)
+            )
+            if expired:
+                self._fail_session(
+                    UnavailableError("queued lease reached its local deadline"),
+                    epoch=assignment.session_epoch,
+                    session=assignment.session,
+                )
+                return
+            self._start_handler(assignment)
+
+    def _start_handler(self, assignment: _Assignment) -> None:
+        assignment.started = True
+        if self._controller is None:
+            self._fixed_running += 1
+        task = asyncio.create_task(
+            self._handle_assignment(assignment),
+            name=f"ringo-handler-{assignment.task.id}-{assignment.attempt}",
+        )
+        assignment.handler_task = task
+        self._handler_tasks.add(task)
+        task.add_done_callback(
+            lambda completed: self._handler_done(assignment, completed)
+        )
+
+    def _handler_done(self, assignment: _Assignment, task: asyncio.Task[Any]) -> None:
+        _consume_task_exception(task)
+        self._handler_tasks.discard(task)
+        if task.cancelled() and assignment.alive:
+            # Task.cancel before its first step never enters the coroutine's
+            # cancellation handler. Stop that handlerless lease synchronously;
+            # the existing release owner still accounts for actual Task.done.
+            self._fail_session(
+                UnavailableError("task handler cancelled before settlement"),
+                epoch=assignment.session_epoch,
+                session=assignment.session,
+            )
+            self._invalidate_assignment(assignment)
+            if self._inflight.get(assignment.identity) is assignment:
+                self._inflight.pop(assignment.identity)
+            renewal = assignment.renew_task
+            if renewal is not None and not renewal.done():
+                renewal.cancel()
+        self._schedule_capacity_release(assignment)
+
+    def _schedule_capacity_release(self, assignment: _Assignment) -> None:
+        task = asyncio.create_task(self._release_finished_capacity(assignment))
+        self._capacity_tasks.add(task)
+        task.add_done_callback(self._capacity_tasks.discard)
+        task.add_done_callback(_consume_task_exception)
+
+    async def _release_finished_capacity(self, assignment: _Assignment) -> None:
+        epoch, session = self._session_epoch, self._session_done
+        try:
+            renewal = assignment.renew_task
+            if not assignment.alive and renewal is not None and not renewal.done():
+                # A done callback can cancel renewal before the handler ever
+                # ran. Join that SDK task without releasing another assignment.
+                await asyncio.gather(renewal, return_exceptions=True)
+            await self._release_capacity_if_settled(assignment)
+        except asyncio.CancelledError:
+            raise
+        except Exception as exc:
+            # A delayed release from an old session must not tear down its
+            # replacement. Publishers also fence their actual command epoch.
+            self._fail_session(exc, epoch=epoch, session=session)
+
+    def _fail_session(
+        self,
+        error: BaseException,
+        *,
+        epoch: int | None = None,
+        session: asyncio.Event | None = None,
+    ) -> None:
+        source = self._source
+        if source is None:
+            return
+        if epoch is not None and (epoch != source.epoch or session is not source.done):
+            return
+        self._end_source(source, error)
+
+    async def _handle_assignment(self, assignment: _Assignment) -> None:
+        try:
+            context = assignment.context
+            assert context is not None
+            if not assignment.alive:
+                return
+            if self._force_cancel.is_set() or (
+                assignment.session is not None and assignment.session.is_set()
+            ):
+                # Reader bursts can schedule a handler and observe EOF before
+                # its first step. The immutable origin Event, not a replacement
+                # source or close_started, gates entry into application code.
+                await self._release_assignment(assignment, cancel_context=True)
+                return
+            # One disabled-path check; no monitor task, clock read or RPC for
+            # ordinary workers. The enabled helper enters immediately before
+            # the real handler call and exits before any terminal receipt.
+            watchdog = self._watchdog
+            if watchdog is None:
+                result = self._handler(assignment.task, context)
+                if inspect.isawaitable(result):
+                    await result
+            else:
+                await watchdog.invoke(assignment, context)
+            if assignment.alive and not self._force_cancel.is_set():
+                await self._finish(assignment, "ack")
+        except _WatchdogFenced:
+            # A deadline racing graceful close abandons only its own stopped
+            # invocation; other handlers keep their live drain/ACK/renewal lane.
+            await self._release_assignment(assignment, cancel_context=True)
+        except asyncio.CancelledError:
+            # Cancellation is not a NACK, but an application-cancelled handler
+            # must not retain a live, perpetually renewed lease and capacity.
+            if assignment.alive:
+                self._fail_session(
+                    UnavailableError("task handler cancelled before settlement"),
+                    epoch=assignment.session_epoch,
+                    session=assignment.session,
+                )
+                await self._release_assignment(assignment, cancel_context=True)
+            raise
+        except RetryTaskError as exc:
+            if assignment.alive and not self._force_cancel.is_set():
+                await self._finish(assignment, "nack", exc)
+        except PermanentTaskError as exc:
+            if assignment.alive and not self._force_cancel.is_set():
+                await self._finish(assignment, "reject", exc)
+        except Exception as exc:
+            if assignment.alive and not self._force_cancel.is_set():
+                await self._finish(assignment, "nack", exc)
+        # Capacity belongs to the complete handler task, including receipt
+        # consumption; its done callback releases the slot exactly once.
+
+    def _watchdog_unresponsive(self, assignment: _Assignment, reason: str) -> None:
+        """Fence before notifying: never NACK while old business still runs."""
+        self._invalidate_assignment(assignment)
+        renewal = assignment.renew_task
+        if renewal is not None:
+            renewal.cancel()
+        # Do not recycle local capacity; the real handler's done callback owns
+        # it even if run() returns while the application swallows cancellation.
+        self._close_started.set()
+        self._publisher_closed = True
+        source = self._source
+        if source is not None:
+            # This is an intentional Worker-wide hard fence, not a stale RPC
+            # result: even a replacement source must stop admitting new work.
+            self._end_source(
+                source, UnavailableError("unresponsive timed-out handler; Worker fenced")
+            )
+        detail = {
+            "worker_id": self._worker_id,
+            "queue": self._queue.name,
+            "task_id": assignment.task.id,
+            "attempt": assignment.attempt,
+            "reason": reason,
+        }
+        if self._runtime is not None:
+            self._runtime.request_restart("task cancellation grace exceeded", detail)
+        else:
+            self._watchdog_error = RuntimeError(
+                "task cancellation grace exceeded: Worker stopped accepting tasks; "
+                "a handler is still running; use process supervision for hard isolation"
+            )
+
+    async def _finish(
+        self,
+        assignment: _Assignment,
+        operation: str,
+        error: BaseException | None = None,
+    ) -> None:
+        if not assignment.alive or assignment.completion_started:
+            return
+        assignment.completion_started = True
+        request: queue_pb2.WorkRequest
+        text = _safe_error(error) if error is not None else ""
+        if operation == "ack":
+            request = queue_pb2.WorkRequest(
+                ack=queue_pb2.AckRequest(lease=assignment.ref)
+            )
+        elif operation == "nack":
+            nack = queue_pb2.NackRequest(lease=assignment.ref, error=text)
+            if isinstance(error, RetryTaskError) and error.delay is not None:
+                if (
+                    not isinstance(error.delay, (int, float))
+                    or not math.isfinite(error.delay)
+                    or error.delay < 0
+                ):
+                    # Treat an invalid override as an ordinary failure rather
+                    # than allowing a negative Duration onto the wire.
+                    nack.error = _safe_error(error)
+                else:
+                    try:
+                        nack.retry_delay_override.CopyFrom(
+                            _pb.td_to_dur(timedelta(seconds=error.delay))
+                        )
+                    except Exception:
+                        # An unrepresentable override is treated as a normal
+                        # NACK; never strand the assignment on serialization.
+                        nack.error = _safe_error(error)
+            request = queue_pb2.WorkRequest(nack=nack)
+        else:
+            request = queue_pb2.WorkRequest(
+                reject=queue_pb2.RejectRequest(lease=assignment.ref, error=text)
+            )
+        try:
+            # Task.cancel() can be pending after a handler returns/throws without
+            # suspending. Queue.put() on the unbounded outbox is a fast path, so
+            # deliver cancellation before any terminal frame reaches the writer.
+            # Keep this inside our cancellation cleanup, including NACK/reject
+            # callers in handler except blocks. A consumed cancellation remains
+            # a valid normal result; its Task.cancelling() count is not a gate.
+            await asyncio.sleep(0)
+            if assignment.watchdog_expired and (
+                not assignment.alive
+                or self._close_started.is_set()
+                or self._force_cancel.is_set()
+            ):
+                await self._release_assignment(assignment, cancel_context=True)
+                return
+            result = await self._send_command(
+                request,
+                expected_epoch=assignment.session_epoch,
+                expected_session=assignment.session,
+            )
+        except asyncio.CancelledError:
+            # This await also runs inside the handler's except blocks (NACK
+            # and reject), where a sibling cancellation except cannot catch it.
+            if assignment.alive:
+                self._fail_session(
+                    UnavailableError("task settlement cancelled before receipt"),
+                    epoch=assignment.session_epoch,
+                    session=assignment.session,
+                )
+                await self._release_assignment(assignment, cancel_context=True)
+            raise
+        except Exception as exc:
+            # The outcome may have committed even without a receipt. Never
+            # replay it or keep renewing a lease with no live handler forever.
+            self._fail_session(
+                exc, epoch=assignment.session_epoch, session=assignment.session
+            )
+            await self._release_assignment(assignment, cancel_context=True)
+            return
+        detail = result.error if result.HasField("error") else None
+        if detail is not None and detail.code == queue_pb2.ERROR_CODE_LEASE_LOST:
+            await self._lease_lost(assignment)
+            return
+        if detail is not None and detail.code != queue_pb2.ERROR_CODE_UNSPECIFIED:
+            self._fail_session(
+                error_from_detail(detail),
+                epoch=assignment.session_epoch,
+                session=assignment.session,
+            )
+            await self._release_assignment(assignment, cancel_context=True)
+            return
+        congestion = isinstance(error, RetryTaskError) and error.congestion
+        await self._release_assignment(
+            assignment,
+            success=operation == "ack",
+            congestion=congestion,
+        )
+
+    async def _release_assignment(
+        self,
+        assignment: _Assignment,
+        *,
+        cancel_context: bool = False,
+        success: bool = False,
+        congestion: bool = False,
+    ) -> None:
+        if not assignment.alive:
+            return
+        assignment.alive = False
+        assignment.completion_success = success
+        assignment.completion_congestion = congestion
+        if self._inflight.get(assignment.identity) is assignment:
+            self._inflight.pop(assignment.identity)
+        if not assignment.started:
+            self._waiting_assignments = deque(
+                pending
+                for pending in self._waiting_assignments
+                if pending is not assignment
+            )
+        if cancel_context and assignment.context is not None:
+            assignment.context._cancel()
+        if (
+            assignment.renew_task is not None
+            and assignment.renew_task is not asyncio.current_task()
+        ):
+            assignment.renew_task.cancel()
+            await asyncio.gather(assignment.renew_task, return_exceptions=True)
+        await self._release_capacity_if_settled(assignment)
+
+    async def _lease_lost(self, assignment: _Assignment) -> None:
+        if not assignment.alive:
+            return
+        await self._release_assignment(assignment, cancel_context=True)
+        handler = assignment.handler_task
+        if (
+            handler is not None
+            and handler is not asyncio.current_task()
+            and not handler.done()
+        ):
+            # Do not await an application coroutine that may intentionally
+            # swallow cancellation. Its capacity is released by its finally.
+            handler.cancel()
+        await self._release_capacity_if_settled(assignment)
+
+    async def _renew_loop(self, assignment: _Assignment) -> None:
+        loop = asyncio.get_running_loop()
+        if assignment.renew_deadline is None:
+            assignment.renew_deadline = (
+                loop.time()
+                + (assignment.until - datetime.now(timezone.utc)).total_seconds()
+            )
+        while assignment.alive and not self._force_cancel.is_set():
+            remaining = assignment.renew_deadline - loop.time()
+            delay = max(0.01, remaining - self._lease_duration.total_seconds() / 3)
+            try:
+                await asyncio.wait_for(self._force_cancel.wait(), timeout=delay)
+                return
+            except asyncio.TimeoutError:
+                pass
+            if not assignment.alive or self._force_cancel.is_set():
+                return
+            request = queue_pb2.WorkRequest(
+                extend_lease=queue_pb2.ExtendLeaseRequest(
+                    lease=assignment.ref,
+                    extension=_pb.td_to_dur(self._lease_duration),
+                )
+            )
+            sent_at = loop.time()
+            sent_wall = datetime.now(timezone.utc)
+            try:
+                result = await self._send_command(
+                    request,
+                    expected_epoch=assignment.session_epoch,
+                    expected_session=assignment.session,
+                )
+            except asyncio.CancelledError:
+                raise
+            except Exception as exc:
+                self._fail_session(
+                    exc, epoch=assignment.session_epoch, session=assignment.session
+                )
+                await self._lease_lost(assignment)
+                return
+            detail = result.error if result.HasField("error") else None
+            if detail is not None and detail.code == queue_pb2.ERROR_CODE_LEASE_LOST:
+                await self._lease_lost(assignment)
+                return
+            if detail is not None and detail.code != queue_pb2.ERROR_CODE_UNSPECIFIED:
+                self._fail_session(
+                    error_from_detail(detail),
+                    epoch=assignment.session_epoch,
+                    session=assignment.session,
+                )
+                await self._lease_lost(assignment)
+                return
+            if (
+                not assignment.alive
+                or self._force_cancel.is_set()
+                or assignment.session_epoch != self._session_epoch
+                or assignment.session is not self._session_done
+                or assignment.session is None
+                or assignment.session.is_set()
+            ):
+                return
+            # WorkResult has no renewed deadline. Server processing cannot
+            # precede publication, so send-time + extension is a conservative
+            # lower bound, including queued writes/backpressure. Receipt time
+            # must not grant an extra lease. Monotonic scheduling survives wall
+            # clock changes; keep the private wall-clock field conservative too.
+            assignment.renew_deadline = sent_at + self._lease_duration.total_seconds()
+            assignment.until = sent_wall + self._lease_duration
+
+    async def _report_progress(
+        self, assignment: _Assignment, current: float, total: float, message: str
+    ) -> None:
+        if not assignment.alive:
+            raise LeaseLostError("task lease is no longer valid")
+        request = queue_pb2.WorkRequest(
+            progress=queue_pb2.ReportProgressRequest(
+                lease=assignment.ref,
+                current=current,
+                total=total,
+                message=_safe_error(message),
+            )
+        )
+        result = await self._send_command(
+            request,
+            expected_epoch=assignment.session_epoch,
+            expected_session=assignment.session,
+        )
+        detail = result.error if result.HasField("error") else None
+        if detail is not None and detail.code == queue_pb2.ERROR_CODE_LEASE_LOST:
+            await self._lease_lost(assignment)
+            raise LeaseLostError(detail.message)
+        if detail is not None and detail.code != queue_pb2.ERROR_CODE_UNSPECIFIED:
+            raise error_from_detail(detail)
+
+    async def _publish_fixed_capacity(self) -> None:
+        epoch, session = self._session_epoch, self._session_done
+        if (
+            self._publisher_closed
+            or self._close_started.is_set()
+            or session is None
+            or session.is_set()
+        ):
+            return
+        self._fixed_capacity_generation += 1
+        generation = self._fixed_capacity_generation
+        task = self._fixed_publication_task
+        new_publisher = (
+            task is None
+            or task.done()
+            or self._fixed_publication_source != (epoch, session)
+        )
+        if new_publisher:
+            self._fixed_publication_waiters = []
+            self._fixed_publication_source = (epoch, session)
+        waiters = self._fixed_publication_waiters
+        accepted: asyncio.Future[None] = asyncio.get_running_loop().create_future()
+        accepted.add_done_callback(_consume_task_exception)
+        waiters.append((generation, accepted))
+        if new_publisher:
+            task = asyncio.create_task(
+                self._drain_fixed_capacity(epoch, session, waiters),
+                name=f"ringo-credit-{self._worker_id}-{epoch}",
+            )
+            self._fixed_publication_task = task
+            assert self._source is not None
+            self._source.publisher = task
+            self._capacity_tasks.add(task)
+            task.add_done_callback(
+                lambda completed: self._fixed_publication_done(completed, waiters)
+            )
+        # A resize waits for its generation's accepted snapshot, not for the
+        # publisher to become idle: a continuously busy stream may never idle.
+        # Caller/grace cancellation cannot abandon the shared publication.
+        await asyncio.shield(accepted)
+
+    def _fixed_publication_done(
+        self,
+        task: asyncio.Task[None],
+        waiters: list[tuple[int, asyncio.Future[None]]],
+    ) -> None:
+        self._capacity_tasks.discard(task)
+        error = asyncio.CancelledError() if task.cancelled() else task.exception()
+        # Task.cancel() before the coroutine's first step never enters its
+        # finally. Complete this source's waiters even at that boundary.
+        for _, future in waiters:
+            if not future.done():
+                if error is None:
+                    future.set_result(None)
+                else:
+                    future.set_exception(error)
+        waiters.clear()
+
+    async def _drain_fixed_capacity(
+        self,
+        epoch: int,
+        session: asyncio.Event,
+        waiters: list[tuple[int, asyncio.Future[None]]],
+    ) -> None:
+        error: BaseException | None = None
+        try:
+            async with self._publisher_lock:
+                while True:
+                    for attempt in range(20):
+                        if (
+                            self._publisher_closed
+                            or self._close_started.is_set()
+                            or epoch != self._session_epoch
+                            or session is not self._session_done
+                            or session.is_set()
+                        ):
+                            return
+                        # Absolute credit is consumed by Assign frames. Coalesce
+                        # requests, never equal numeric values; an in-transit
+                        # assignment dirties the generation even if a completion
+                        # leaves occupancy unchanged. Retry re-samples all state.
+                        generation = self._fixed_capacity_generation
+                        available = max(
+                            0,
+                            self._concurrency
+                            - self._fixed_running
+                            - len(self._waiting_assignments),
+                        )
+                        update = queue_pb2.CreditUpdate(available=available)
+                        if self._fixed_dynamic_target:
+                            update.target_concurrency = self._concurrency
+                        result = await self._send_command(
+                            queue_pb2.WorkRequest(credit=update),
+                            expected_epoch=epoch,
+                            expected_session=session,
+                        )
+                        if (
+                            not result.HasField("error")
+                            or result.error.code == queue_pb2.ERROR_CODE_UNSPECIFIED
+                        ):
+                            source = self._source
+                            if source is not None and source.done is session:
+                                source.accepted_credit = available
+                                source.credit_generation = generation
+                            break
+                        if (
+                            result.error.metadata.get("reason") != "credit_capacity"
+                            or attempt == 19
+                        ):
+                            raise error_from_detail(result.error)
+                        await asyncio.sleep(0.01)
+                    remaining = []
+                    for requested, future in waiters:
+                        if requested <= generation:
+                            if not future.done():
+                                future.set_result(None)
+                        else:
+                            remaining.append((requested, future))
+                    waiters[:] = remaining
+                    if generation == self._fixed_capacity_generation:
+                        return
+        except BaseException as exc:
+            error = exc
+            if not isinstance(exc, asyncio.CancelledError):
+                self._fail_session(exc, epoch=epoch, session=session)
+            raise
+        finally:
+            # Session replacement/close releases only this publisher's waiters.
+            # A late old failure never completes or cancels a replacement batch.
+            for _, future in waiters:
+                if not future.done():
+                    if error is None:
+                        future.set_result(None)
+                    else:
+                        future.set_exception(error)
+            waiters.clear()
+
+    async def _publish_controller_initial(self, capacity: int) -> None:
+        async with self._publisher_lock:
+            if (
+                self._publisher_closed
+                or self._close_started.is_set()
+                or self._session_done is None
+                or self._session_done.is_set()
+            ):
+                raise UnavailableError("Work stream is not connected")
+            ok = await self._publish_dynamic_locked(capacity, 0)
+            if not ok:
+                raise UnavailableError("initial dynamic credit was not accepted")
+
+    async def _publish_controller_probe(
+        self, member: _Member, capacity: int, generation: int
+    ) -> None:
+        async with self._publisher_lock:
+            if (
+                self._publisher_closed
+                or self._close_started.is_set()
+                or self._member is not member
+                or not member.probe
+                or member.probe_generation != generation
+                or self._session_done is None
+                or self._session_done.is_set()
+            ):
+                return
+            await self._publish_dynamic_locked(
+                capacity,
+                1,
+                probe_member=member,
+                probe_generation=generation,
+            )
+
+    async def _publish_controller_target(self, capacity: int) -> bool:
+        """Publish target/zero credit for a live shared-controller stream."""
+        async with self._publisher_lock:
+            if self._publisher_closed:
+                return True
+            if self._session_done is None or self._session_done.is_set():
+                return True
+            return await self._publish_dynamic_locked(capacity, 0)
+
+    async def _withdraw_controller_probe(
+        self,
+        member: _Member,
+        capacity: int,
+        *,
+        force_target: bool = False,
+    ) -> bool:
+        async with self._publisher_lock:
+            if self._publisher_closed:
+                return True
+            if self._session_done is None or self._session_done.is_set():
+                # The old stream is definitively gone; its credit no longer
+                # exists server-side and the controller may reclaim the token.
+                return True
+            return await self._publish_dynamic_locked(
+                capacity,
+                0,
+                allow_closing=True,
+                explicit_target=force_target,
+            )
+
+    async def _publish_dynamic_locked(
+        self,
+        target: int,
+        available: int,
+        *,
+        allow_closing: bool = False,
+        explicit_target: bool = False,
+        probe_member: _Member | None = None,
+        probe_generation: int = 0,
+    ) -> bool:
+        controller = self._controller
+        assert controller is not None
+        epoch, session = self._session_epoch, self._session_done
+        for attempt in range(20):
+            if (
+                self._publisher_closed
+                or (self._close_started.is_set() and not allow_closing)
+                or epoch != self._session_epoch
+                or self._session_done is None
+                or self._session_done.is_set()
+            ):
+                return self._session_done is None or self._session_done.is_set()
+            if probe_member is not None and (
+                not probe_member.probe
+                or probe_member.probe_generation != probe_generation
+            ):
+                # The assignment consumed this probe while a previous
+                # capacity update was in flight; never replay available=1.
+                return True
+            # target and available are one atomic protocol snapshot.  A
+            # credit_capacity rejection must be retried with a fresh target /
+            # in-flight sample, never by replaying the old snapshot.
+            current_target = controller.capacity if not explicit_target else target
+            update = queue_pb2.CreditUpdate(
+                available=max(0, available),
+                target_concurrency=current_target,
+            )
+            result = await self._send_command(
+                queue_pb2.WorkRequest(credit=update),
+                expected_epoch=epoch,
+                expected_session=session,
+            )
+            if (
+                not result.HasField("error")
+                or result.error.code == queue_pb2.ERROR_CODE_UNSPECIFIED
+            ):
+                source = self._source
+                if source is not None and source.done is session:
+                    source.accepted_credit = available
+                    source.credit_generation = probe_generation
+                return True
+            if (
+                result.error.metadata.get("reason") != "credit_capacity"
+                or attempt == 19
+            ):
+                error = error_from_detail(result.error)
+                self._fail_session(error, epoch=epoch, session=session)
+                raise error
+            await asyncio.sleep(0.01)
+        return False
+
+    async def _send_command(
+        self,
+        request: queue_pb2.WorkRequest,
+        *,
+        expected_epoch: int | None = None,
+        expected_session: asyncio.Event | None = None,
+    ) -> queue_pb2.WorkResult:
+        source = self._source
+        if source is None or source.done.is_set():
+            raise UnavailableError("Work stream is not connected")
+        if expected_epoch is not None and (
+            expected_epoch != source.epoch or expected_session is not source.done
+        ):
+            raise UnavailableError("command belongs to a replaced Work session")
+        if source.next_request_id > _MAX_REQUEST_ID:
+            error = UnavailableError(
+                "Work request ID space exhausted; reconnect required"
+            )
+            self._end_source(source, error)
+            raise error
+        # This is the actual FIFO publication point, shared by every producer
+        # including drain. No await separates source validation, numbering,
+        # pending registration and insertion. Even an immediate receipt can
+        # find its Future, and cancellation never rolls a consumed ID back.
+        rid = source.next_request_id
+        source.next_request_id += 1
+        request.request_id = rid
+        future: asyncio.Future[queue_pb2.WorkResult] = (
+            asyncio.get_running_loop().create_future()
+        )
+        source.pending[rid] = future
+        try:
+            self._publish_frame(source, request)
+            async with asyncio.timeout(self._command_timeout):
+                return await future
+        except asyncio.CancelledError:
+            if future.done() and not future.cancelled():
+                future.exception()
+            raise
+        except Exception as exc:
+            # Source ownership, rather than the current Work call, fences a
+            # timeout or publication failure. Unknown mutations are not replayed.
+            self._end_source(source, exc)
+            if future.done() and not future.cancelled():
+                # Publication can fail before we ever await this Future. End
+                # the source first, then consume its newly installed exception.
+                future.exception()
+            raise
+        finally:
+            if source.pending.get(rid) is future:
+                source.pending.pop(rid)
+
+
+__all__ = ["TaskContext", "Worker"]