ringo-task-queue 0.1.0.dev0__py3-none-win_amd64.whl

ringo_task_queue/client.py

@@ -0,0 +1,863 @@
+"""Connection management and the user-facing client/queue APIs."""
+
+from __future__ import annotations
+
+import asyncio
+import os
+import secrets
+from datetime import timedelta
+from typing import TYPE_CHECKING, Any, Awaitable, Callable, Coroutine, Iterable, Literal
+
+import grpc
+from grpc_health.v1 import health_pb2, health_pb2_grpc
+
+from . import _pb
+from ._proto.ringo.v1 import queue_pb2, queue_pb2_grpc
+from ._version import PROTOCOL_MAJOR, PROTOCOL_MINOR, SDK_VERSION
+from .errors import (
+    IncompatibleVersionError,
+    InvalidArgumentError,
+    RingoError,
+    UnavailableError,
+    map_rpc_error,
+)
+from .models import (
+    BatchResult,
+    EnqueueResult,
+    PostgresStorage,
+    SQLiteStorage,
+    Task,
+    TaskSpec,
+)
+
+if TYPE_CHECKING:  # pragma: no cover
+    from .daemon import LocalDaemon
+    from .worker import Worker
+
+SDK_NAME = "ringo-python"
+
+DEFAULT_RPC_TIMEOUT = 30.0
+_MAX_ENQUEUE_BATCH = 1000
+_RECONNECT_ATTEMPTS = 3
+_BACKOFF_SECONDS = (0.2, 0.6, 1.5)
+
+#: Supported values for the remote ``compression`` option: ``"none"``
+#: (the default) or ``"gzip"``. ``None`` means "not specified": the channel
+#: default applies and a legacy raw ``channel_options`` compression entry
+#: keeps its meaning. The runtime validation in
+#: :func:`_normalize_compression` still rejects any other value for dynamic
+#: callers.
+CompressionOption = Literal["none", "gzip"]
+_COMPRESSION_CHOICES: tuple[str, ...] = ("none", "gzip")
+#: gRPC channel option that controls the same thing as ``compression``. An
+#: explicit ``compression`` (even ``"none"``) combined with this raw option is
+#: an ambiguity error, mirroring the Node SDK.
+_COMPRESSION_CHANNEL_OPTION = "grpc.default_compression_algorithm"
+
+
+def _normalize_compression(value: CompressionOption | None) -> grpc.Compression | None:
+    """Map the public ``compression`` option onto ``grpc.Compression``.
+
+    ``None`` (unspecified) maps to ``None`` so the channel default applies.
+    Anything else outside ``_COMPRESSION_CHOICES`` is a clear
+    ``InvalidArgumentError`` raised before any network activity.
+    """
+    if value is None:
+        return None
+    if not isinstance(value, str) or value not in _COMPRESSION_CHOICES:
+        raise InvalidArgumentError(
+            f"unsupported compression {value!r}; expected one of"
+            f" {', '.join(_COMPRESSION_CHOICES)}"
+        )
+    if value == "gzip":
+        return grpc.Compression.Gzip
+    return grpc.Compression.NoCompression
+
+
+def _check_compression_channel_option_conflict(
+    channel_options: Iterable[tuple[str, Any]] | None,
+    compression: grpc.Compression | None,
+) -> None:
+    """Reject an explicit ``compression`` alongside the raw gRPC option.
+
+    When ``compression`` is unspecified (``None``), raw ``channel_options``
+    keep their legacy meaning for advanced users.
+    """
+    if compression is None:
+        return
+    for key, _ in channel_options or ():
+        if key == _COMPRESSION_CHANNEL_OPTION:
+            raise InvalidArgumentError(
+                f"compression and channel_options[{_COMPRESSION_CHANNEL_OPTION!r}]"
+                " must not both be set"
+            )
+
+
+def _consume_task_exception(task: asyncio.Task[Any]) -> None:
+    if not task.cancelled():
+        task.exception()
+
+
+def _parse_endpoint(endpoint: str) -> tuple[str, bool]:
+    """Return ``(authority, use_tls)`` for an endpoint string."""
+    if not endpoint:
+        raise InvalidArgumentError("endpoint is required")
+    use_tls = False
+    rest = endpoint
+    for scheme, tls in (
+        ("grpc://", False),
+        ("http://", False),
+        ("grpcs://", True),
+        ("https://", True),
+    ):
+        if rest.startswith(scheme):
+            rest = rest[len(scheme) :]
+            use_tls = tls
+            break
+    rest = rest.rstrip("/")
+    if ":" not in rest:
+        rest = f"{rest}:7233"
+    return rest, use_tls
+
+
+#: Bounded automatic restarts of a crashed local daemon (design §31).
+_MAX_DAEMON_RESTARTS = 3
+
+
+class _Connection:
+    """Owns the grpc.aio channel, negotiation, and retried unary calls."""
+
+    def __init__(
+        self,
+        endpoint: str,
+        token: str | None = None,
+        *,
+        channel_options: Iterable[tuple[str, Any]] | None = None,
+        compression: grpc.Compression | None = None,
+    ) -> None:
+        self._endpoint, self._tls = _parse_endpoint(endpoint)
+        self._token = token
+        self._channel_options = list(channel_options or ())
+        _check_compression_channel_option_conflict(
+            self._channel_options, compression
+        )
+        # Channel-level default compression; survives every reconnect because
+        # _open_channel() re-applies it to each new aio channel. ``None`` means
+        # unspecified, so legacy raw channel options decide the encoding.
+        self._compression = compression
+        self._channel: grpc.aio.Channel | None = None
+        self._stub: queue_pb2_grpc.QueueServiceStub | None = None
+        self._session_minor = PROTOCOL_MINOR
+        self._session_token = ""
+        self._lock = asyncio.Lock()
+        self._closed = False
+        #: Optional client-provided reconnect hook (daemon restart + channel
+        #: rebuild). Invoked under ``_lock`` instead of plain ``negotiate``.
+        self.on_reconnect: Callable[[], Awaitable[None]] | None = None
+
+    # -- lifecycle ---------------------------------------------------------
+
+    def _open_channel(self) -> None:
+        if self._tls:
+            credentials = grpc.ssl_channel_credentials()
+            self._channel = grpc.aio.secure_channel(
+                self._endpoint,
+                credentials,
+                options=self._channel_options,
+                compression=self._compression,
+            )
+        else:
+            self._channel = grpc.aio.insecure_channel(
+                self._endpoint,
+                options=self._channel_options,
+                compression=self._compression,
+            )
+        self._stub = queue_pb2_grpc.QueueServiceStub(self._channel)
+
+    async def connect(self, *, check_health: bool = False) -> None:
+        self._open_channel()
+        if check_health:
+            await self._await_serving()
+        await self.negotiate()
+
+    async def _await_serving(self, timeout: float = 5.0) -> None:
+        """Wait until the health service reports SERVING (local daemon
+        readiness is announced on stdout before the gRPC health state may be
+        observable, so poll briefly instead of assuming)."""
+        assert self._channel is not None
+        health_stub = health_pb2_grpc.HealthStub(self._channel)
+        deadline = asyncio.get_running_loop().time() + timeout
+        last: Exception | None = None
+        while True:
+            try:
+                response = await health_stub.Check(
+                    health_pb2.HealthCheckRequest(),
+                    timeout=2.0,
+                    metadata=self._metadata(),
+                )
+                if response.status == health_pb2.HealthCheckResponse.SERVING:
+                    return
+                last = RingoError("daemon health status is not SERVING")
+            except grpc.aio.AioRpcError as exc:
+                last = exc
+            if asyncio.get_running_loop().time() >= deadline:
+                raise RingoError(f"daemon did not report SERVING health: {last}")
+            await asyncio.sleep(0.05)
+
+    async def reset(self, endpoint: str) -> None:
+        """Point the connection at a new endpoint (local daemon restart)."""
+        self._endpoint, self._tls = _parse_endpoint(endpoint)
+        channel, self._channel = self._channel, None
+        self._stub = None
+        if channel is not None:
+            await channel.close()
+        self._open_channel()
+
+    async def negotiate(self) -> None:
+        """Perform the major=1/minor=0 version handshake (ADR 0002)."""
+        stub = self._require_stub()
+        request = queue_pb2.NegotiateRequest(
+            sdk_name=SDK_NAME,
+            sdk_version=SDK_VERSION,
+            protocol_major=PROTOCOL_MAJOR,
+            protocol_minor=PROTOCOL_MINOR,
+        )
+        try:
+            response = await stub.Negotiate(
+                request, timeout=DEFAULT_RPC_TIMEOUT, metadata=self._metadata()
+            )
+        except grpc.aio.AioRpcError as exc:
+            raise map_rpc_error(exc) from exc
+        if response.protocol_major != PROTOCOL_MAJOR:
+            raise IncompatibleVersionError(
+                f"server speaks protocol major {response.protocol_major},"
+                f" SDK requires major {PROTOCOL_MAJOR}; upgrade one side"
+            )
+        self._session_minor = min(response.protocol_minor, PROTOCOL_MINOR)
+        self._session_token = response.session_token
+
+    async def close(self) -> None:
+        self._closed = True
+        channel, self._channel = self._channel, None
+        self._stub = None
+        if channel is not None:
+            await channel.close()
+
+    # -- helpers -----------------------------------------------------------
+
+    def _require_stub(self) -> queue_pb2_grpc.QueueServiceStub:
+        if self._stub is None:
+            raise RingoError("connection is closed")
+        return self._stub
+
+    def _metadata(self) -> list[tuple[str, str]]:
+        if self._token:
+            return [("authorization", f"Bearer {self._token}")]
+        return []
+
+    def request_context(self) -> queue_pb2.RequestContext:
+        return queue_pb2.RequestContext(
+            protocol_major=PROTOCOL_MAJOR,
+            protocol_minor=self._session_minor,
+            session_token=self._session_token,
+        )
+
+    @property
+    def endpoint(self) -> str:
+        return self._endpoint
+
+    async def call(
+        self,
+        name: str,
+        request: Any,
+        *,
+        timeout: float = DEFAULT_RPC_TIMEOUT,
+        retry_safe: bool = False,
+        context_attr: str | None = "context",
+    ) -> Any:
+        """Issue a single unary RPC.
+
+        Submission is never silently replayed: an ``UNAVAILABLE`` may mean the
+        server committed the operation (a lost ACK response, a claimed lease)
+        so the original error is raised after best-effort recovery (daemon
+        restart + fresh negotiation) prepares the connection for *future*
+        calls. Only genuinely read-only RPCs (``GetTask``) pass
+        ``retry_safe=True`` and may be re-sent, each time with a freshly
+        negotiated ``RequestContext``.
+        """
+        rpc, sent_token = await self._begin_unary(
+            name, request, timeout=timeout, context_attr=context_attr
+        )
+        cause: BaseException | None = None
+        err: RingoError | None = None
+        try:
+            return await rpc
+        except grpc.aio.AioRpcError as exc:
+            cause = exc
+            err = map_rpc_error(exc)
+        except asyncio.CancelledError as exc:
+            # grpc.aio raises CancelledError when a reconnect closes its old
+            # channel. Preserve genuine caller cancellation, but surface this
+            # transport race as UNAVAILABLE.
+            task = asyncio.current_task()
+            if task is not None and task.cancelling():
+                raise
+            cause = exc
+            err = UnavailableError("RPC interrupted by connection recovery")
+
+        assert err is not None and cause is not None
+        if self._closed:
+            raise err from cause
+        if isinstance(err, IncompatibleVersionError):
+            # Session rejection happens in authorization before the operation
+            # can execute. It is therefore safe to negotiate a fresh short-lived
+            # token and retry once, including ordinary token expiry with no
+            # concurrent daemon restart.
+            await self._recover_best_effort()
+            async with self._lock:
+                rotated = self._session_token != sent_token
+            if rotated:
+                rpc, _ = await self._begin_unary(
+                    name, request, timeout=timeout, context_attr=context_attr
+                )
+                try:
+                    return await rpc
+                except grpc.aio.AioRpcError as retry_exc:
+                    raise map_rpc_error(retry_exc) from retry_exc
+            raise err from cause
+        if not isinstance(err, UnavailableError):
+            raise err from cause
+
+        # Recovery prepares the connection for future calls. Mutations are not
+        # replayed because an unavailable response may hide a committed write.
+        await self._recover_best_effort()
+        if not retry_safe:
+            raise err from cause
+        rpc, _ = await self._begin_unary(
+            name, request, timeout=timeout, context_attr=context_attr
+        )
+        try:
+            return await rpc
+        except grpc.aio.AioRpcError as retry_exc:
+            raise map_rpc_error(retry_exc) from retry_exc
+
+    async def _begin_unary(
+        self,
+        name: str,
+        request: Any,
+        *,
+        timeout: float,
+        context_attr: str | None,
+    ) -> tuple[Any, str]:
+        """Snapshot a stub and fresh session context outside reconnect races."""
+        async with self._lock:
+            if context_attr is not None:
+                getattr(request, context_attr).CopyFrom(self.request_context())
+            stub = self._require_stub()
+            token = self._session_token
+            rpc = getattr(stub, name)(
+                request, timeout=timeout, metadata=self._metadata()
+            )
+            return rpc, token
+
+    async def _recover_best_effort(self) -> None:
+        """Bounded reconnect so later calls can succeed; failures here are
+        swallowed because the original call error is what the caller sees."""
+        for attempt in range(_RECONNECT_ATTEMPTS):
+            try:
+                await asyncio.sleep(
+                    _BACKOFF_SECONDS[min(attempt, len(_BACKOFF_SECONDS) - 1)]
+                )
+                async with self._lock:
+                    if self._closed:
+                        return
+                    if self.on_reconnect is not None:
+                        await self.on_reconnect()
+                    else:
+                        await self.negotiate()
+                return
+            except Exception:
+                continue
+
+
+def new_worker_id() -> str:
+    return f"py-{os.getpid()}-{secrets.token_hex(4)}"
+
+
+# ---------------------------------------------------------------------------
+# Public client
+
+
+class RingoClient:
+    """A connected (or connectable) Ringo client.
+
+    Use as an async context manager, or call :meth:`start` / :meth:`close`
+    explicitly::
+
+        async with Ringo.connect("127.0.0.1:7233") as client:
+            queue = client.queue("crawl")
+    """
+
+    def __init__(
+        self,
+        *,
+        endpoint: str | None = None,
+        token: str | None = None,
+        daemon: "LocalDaemon | None" = None,
+        channel_options: Iterable[tuple[str, Any]] | None = None,
+        compression: CompressionOption | None = None,
+    ) -> None:
+        self._endpoint = endpoint
+        self._token = token
+        self._channel_options = channel_options
+        # The value is always validated, but ``compression`` only applies to
+        # remote connections; local mode stays on the (uncompressed) default.
+        normalized = _normalize_compression(compression)
+        if daemon is None:
+            _check_compression_channel_option_conflict(
+                self._channel_options, normalized
+            )
+            self._compression = normalized
+            self._compression_label: CompressionOption = (
+                compression if compression is not None else "none"
+            )
+        else:
+            # Local mode never enables compression; an explicit value is
+            # validated above and then ignored (source compatibility).
+            self._compression = None
+            self._compression_label = "none"
+        self._daemon = daemon
+        # ``_conn`` is built in start(): a local client learns the daemon's
+        # random loopback port only after the daemon announces readiness.
+        self._conn: _Connection | None = None
+        self._queues: dict[str, TaskQueue] = {}
+        self._workers: set[Worker] = set()
+        self._started = False
+        self._closed = False
+        self._start_lock = asyncio.Lock()
+        self._close_lock = asyncio.Lock()
+        self._close_task: asyncio.Task[None] | None = None
+        self._daemon_restart_count = 0
+
+    # -- lifecycle ---------------------------------------------------------
+
+    async def start(self) -> "RingoClient":
+        async with self._start_lock:
+            if self._closed:
+                raise RingoError("client is closed")
+            if not self._started:
+                try:
+                    if self._daemon is not None:
+                        self._endpoint = await self._daemon.start()
+                    self._conn = _Connection(
+                        self._endpoint or "",
+                        self._token,
+                        channel_options=self._channel_options,
+                        compression=self._compression,
+                    )
+                    self._conn.on_reconnect = self._reconnect
+                    await self._conn.connect(check_health=self._daemon is not None)
+                except BaseException:
+                    # Fast cleanup: never leak a half-started local daemon.
+                    await self._stop_daemon()
+                    raise
+                self._started = True
+        return self
+
+    def _require_conn(self) -> _Connection:
+        conn = self._conn
+        if conn is None:
+            raise RingoError(
+                "client is not started; use 'async with' or await client.start()"
+            )
+        return conn
+
+    async def _reconnect(self) -> None:
+        """Reconnect hook: restart a dead local daemon (bounded, ADR 0003 /
+        design §31: at most 3 restarts with backoff), then renegotiate."""
+        conn = self._require_conn()
+        daemon = self._daemon
+        if daemon is not None and daemon.is_dead:
+            if self._daemon_restart_count >= _MAX_DAEMON_RESTARTS:
+                raise UnavailableError(
+                    f"local daemon restart budget ({_MAX_DAEMON_RESTARTS}) exhausted"
+                )
+            self._daemon_restart_count += 1
+            delay = 0.3 * (2 ** (self._daemon_restart_count - 1))
+            await asyncio.sleep(delay)
+            endpoint = await daemon.start()
+            await conn.reset(endpoint)
+        await conn.negotiate()
+
+    async def __aenter__(self) -> "RingoClient":
+        return await self.start()
+
+    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._closed = True
+                self._close_task = asyncio.create_task(
+                    self._close_impl(), name="ringo-client-close"
+                )
+                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:
+        """Attempt every cleanup step and aggregate any failures."""
+        errors: list[Exception] = []
+        # 1. Drain workers (stop taking tasks, wait grace, cancel remainder).
+        for worker in list(self._workers):
+            try:
+                await worker.close()
+            except Exception as exc:  # keep cleaning up the rest
+                errors.append(exc)
+        self._workers.clear()
+        # 2. Close the gRPC channel.
+        conn, self._conn = self._conn, None
+        if conn is not None:
+            try:
+                await conn.close()
+            except Exception as exc:
+                errors.append(exc)
+        # 3. Shut the local daemon down via the lifecycle pipe, then wait.
+        try:
+            await self._stop_daemon()
+        except Exception as exc:
+            errors.append(exc)
+        if errors:
+            raise ExceptionGroup("errors while closing Ringo client", errors)
+
+    async def _stop_daemon(self) -> None:
+        daemon, self._daemon = self._daemon, None
+        if daemon is not None:
+            await daemon.stop()
+
+    async def __aexit__(self, exc_type, exc, tb) -> None:
+        try:
+            await self.close()
+        except ExceptionGroup as cleanup_errors:
+            # A business exception takes precedence; cleanup failures must
+            # not mask it, but should not vanish either.
+            if exc is None:
+                raise
+            for error in cleanup_errors.exceptions:
+                exc.__cause__ = exc.__cause__ or error
+
+    @property
+    def is_local(self) -> bool:
+        return self._daemon is not None
+
+    @property
+    def compression(self) -> CompressionOption:
+        """The effective channel compression: ``"none"`` (default) or ``"gzip"``."""
+        return self._compression_label
+
+    # -- queues ------------------------------------------------------------
+
+    def queue(self, name: str) -> "TaskQueue":
+        if not name:
+            raise InvalidArgumentError("queue name is required")
+        if self._closed:
+            raise RingoError("client is closed")
+        existing = self._queues.get(name)
+        if existing is None:
+            existing = TaskQueue(self, name)
+            self._queues[name] = existing
+        return existing
+
+    def _register_worker(self, worker: "Worker") -> None:
+        self._workers.add(worker)
+
+    def _unregister_worker(self, worker: "Worker") -> None:
+        self._workers.discard(worker)
+
+
+class Ringo:
+    """Factory for local (embedded daemon) and remote clients.
+
+    Both factories are synchronous and return a client whose
+    ``async with``/``await start()`` performs the actual connect and
+    negotiation, so ``async with Ringo.local(...)`` works directly.
+    """
+
+    @staticmethod
+    def local(
+        data_dir: str | os.PathLike[str],
+        storage: str | SQLiteStorage | PostgresStorage = "sqlite",
+        daemon_path: str | os.PathLike[str] | None = None,
+        *,
+        startup_timeout: float = 15.0,
+        stop_timeout: float = 10.0,
+    ) -> RingoClient:
+        """Start a private daemon over SQLite or user-managed PostgreSQL."""
+        if storage == "sqlite" or isinstance(storage, SQLiteStorage):
+            storage_name = "sqlite"
+            postgres_dsn = None
+        elif isinstance(storage, PostgresStorage):
+            storage_name = "postgres"
+            postgres_dsn = storage.dsn
+        else:
+            raise InvalidArgumentError(f"unsupported local storage {storage!r}")
+        # Imported here so remote-only users never pay for subprocess imports.
+        from .daemon import LocalDaemon
+
+        daemon = LocalDaemon(
+            data_dir,
+            storage=storage_name,
+            postgres_dsn=postgres_dsn,
+            daemon_path=daemon_path,
+            startup_timeout=startup_timeout,
+            stop_timeout=stop_timeout,
+        )
+        return RingoClient(daemon=daemon)
+
+    @staticmethod
+    def connect(
+        endpoint: str,
+        token: str | None = None,
+        *,
+        channel_options: Iterable[tuple[str, Any]] | None = None,
+        compression: CompressionOption | None = None,
+    ) -> RingoClient:
+        """Create a client for a remote (or already-running) endpoint.
+
+        ``compression`` is an optional whole-channel setting: ``"none"``
+        (the default) or ``"gzip"``. ``None`` means "not specified", leaving
+        the gRPC channel default (and any raw ``channel_options`` compression
+        entry) in charge. An explicit value combined with the raw channel
+        option ``grpc.default_compression_algorithm`` raises
+        :class:`InvalidArgumentError`; any other value does so as well,
+        immediately and before any network activity.
+        """
+        return RingoClient(
+            endpoint=endpoint,
+            token=token,
+            channel_options=channel_options,
+            compression=compression,
+        )
+
+
+# ---------------------------------------------------------------------------
+# TaskQueue
+
+
+Handler = Callable[[Task, Any], Coroutine[Any, Any, None]]
+
+
+class TaskQueue:
+    """A client bound to one logical queue (``queue + uniqueKey`` dedupe)."""
+
+    def __init__(self, client: RingoClient, name: str) -> None:
+        self._client = client
+        self.name = name
+
+    # -- enqueue ------------------------------------------------------------
+
+    async def enqueue(
+        self,
+        spec: TaskSpec | None = None,
+        /,
+        **kwargs: Any,
+    ) -> EnqueueResult:
+        """Enqueue one task; duplicates return ``created=False`` (no raise)."""
+        if spec is None:
+            spec = TaskSpec(**kwargs)
+        request = queue_pb2.EnqueueRequest(queue=self.name, task=spec.to_proto())
+        request.context.CopyFrom(self._client._require_conn().request_context())
+        response = await self._client._require_conn().call("Enqueue", request)
+        return EnqueueResult.from_proto(response.result)
+
+    async def enqueue_many(self, specs: Iterable[TaskSpec]) -> BatchResult:
+        """Enqueue a batch (max 1000), reporting every item individually."""
+        specs = list(specs)
+        if not specs:
+            raise InvalidArgumentError("enqueue_many requires at least one TaskSpec")
+        if len(specs) > _MAX_ENQUEUE_BATCH:
+            raise InvalidArgumentError(
+                f"enqueue_many accepts at most {_MAX_ENQUEUE_BATCH} tasks"
+            )
+        request = queue_pb2.EnqueueManyRequest(
+            queue=self.name, tasks=[s.to_proto() for s in specs]
+        )
+        request.context.CopyFrom(self._client._require_conn().request_context())
+        response = await self._client._require_conn().call("EnqueueMany", request)
+        return BatchResult(
+            results=[EnqueueResult.from_proto(r) for r in response.results],
+            created_count=response.created_count,
+            duplicate_count=response.duplicate_count,
+        )
+
+    # -- queries ------------------------------------------------------------
+
+    async def get_by_id(self, task_id: str) -> Task:
+        if not task_id:
+            raise InvalidArgumentError("task_id is required")
+        request = queue_pb2.GetTaskRequest(queue=self.name, task_id=task_id)
+        request.context.CopyFrom(self._client._require_conn().request_context())
+        response = await self._client._require_conn().call(
+            "GetTask", request, retry_safe=True
+        )
+        return Task.from_proto(response.task)
+
+    async def get_by_unique_key(self, unique_key: str) -> Task:
+        if not unique_key:
+            raise InvalidArgumentError("unique_key is required")
+        request = queue_pb2.GetTaskRequest(queue=self.name, unique_key=unique_key)
+        request.context.CopyFrom(self._client._require_conn().request_context())
+        response = await self._client._require_conn().call(
+            "GetTask", request, retry_safe=True
+        )
+        return Task.from_proto(response.task)
+
+    async def retry(self, unique_key: str) -> Task:
+        """Reactivate a SUCCEEDED/DEAD task by its business unique key.
+
+        This always resolves through the proto ``unique_key`` field; use
+        :meth:`retry_by_id` for the task-ID form. No string is ever guessed.
+        """
+        if not unique_key:
+            raise InvalidArgumentError("unique_key is required")
+        request = queue_pb2.RetryTaskRequest(queue=self.name)
+        request.unique_key = unique_key
+        request.context.CopyFrom(self._client._require_conn().request_context())
+        response = await self._client._require_conn().call("RetryTask", request)
+        return Task.from_proto(response.task)
+
+    async def retry_by_id(self, task_id: str) -> Task:
+        """Reactivate a SUCCEEDED/DEAD task by its server task ID."""
+        if not task_id:
+            raise InvalidArgumentError("task_id is required")
+        request = queue_pb2.RetryTaskRequest(queue=self.name)
+        request.task_id = task_id
+        request.context.CopyFrom(self._client._require_conn().request_context())
+        response = await self._client._require_conn().call("RetryTask", request)
+        return Task.from_proto(response.task)
+
+    # -- manual claim ---------------------------------------------------------
+
+    async def claim(
+        self,
+        limit: int = 1,
+        *,
+        task_types: Iterable[str] | None = None,
+        wait_timeout: timedelta | None = None,
+        lease_duration: timedelta | None = None,
+        worker_id: str | None = None,
+    ) -> list["LeasedTask"]:
+        """Manually claim tasks (advanced API; not the default worker path)."""
+        if limit < 1 or limit > _MAX_ENQUEUE_BATCH:
+            raise InvalidArgumentError(
+                f"limit must be between 1 and {_MAX_ENQUEUE_BATCH}"
+            )
+        if wait_timeout is not None and wait_timeout < timedelta(0):
+            raise InvalidArgumentError("wait_timeout must not be negative")
+        if lease_duration is not None and lease_duration < timedelta(0):
+            raise InvalidArgumentError("lease_duration must not be negative")
+        request = queue_pb2.ClaimRequest(
+            queue=self.name,
+            worker_id=worker_id or new_worker_id(),
+            task_types=list(task_types or []),
+            limit=limit,
+        )
+        if wait_timeout is not None:
+            request.wait_timeout.CopyFrom(_pb.td_to_dur(wait_timeout))
+        if lease_duration is not None:
+            request.lease_duration.CopyFrom(_pb.td_to_dur(lease_duration))
+        request.context.CopyFrom(self._client._require_conn().request_context())
+        response = await self._client._require_conn().call("Claim", request)
+        return [LeasedTask(self, t) for t in response.tasks]
+
+    # -- worker ---------------------------------------------------------------
+
+    def worker(
+        self,
+        handler: Handler,
+        *,
+        task_types: Iterable[str] | None = None,
+        concurrency: int = 1,
+        lease_duration: timedelta = timedelta(seconds=60),
+        grace_period: timedelta = timedelta(seconds=30),
+    ) -> "Worker":
+        """Create a credit-controlled worker over the frozen Work stream."""
+        from .worker import Worker
+
+        if not callable(handler):
+            raise InvalidArgumentError("handler must be callable")
+        if concurrency < 1 or concurrency > 2**32 - 1:
+            raise InvalidArgumentError("concurrency must fit in a non-zero uint32")
+        if lease_duration < timedelta(0):
+            raise InvalidArgumentError("lease_duration must not be negative")
+        if grace_period < timedelta(0):
+            raise InvalidArgumentError("grace_period must not be negative")
+        worker = Worker(
+            self,
+            handler,
+            task_types=list(task_types or []),
+            concurrency=concurrency,
+            lease_duration=lease_duration,
+            grace_period=grace_period,
+        )
+        self._client._register_worker(worker)
+        return worker
+
+
+# ---------------------------------------------------------------------------
+# Manual lease
+
+
+class LeasedTask:
+    """A task claimed via :meth:`TaskQueue.claim` with its current lease."""
+
+    def __init__(self, queue: TaskQueue, msg: queue_pb2.LeasedTask) -> None:
+        self._queue = queue
+        self.task = Task.from_proto(msg.task)
+        self.attempt: int = msg.lease.attempt
+        self.worker_id: str = msg.lease.worker_id
+        self._token: str = msg.lease.lease_token
+        self._lease_until = _pb.ts_to_dt(msg.lease.until)
+
+    @property
+    def lease_until(self):
+        return self._lease_until
+
+    def _ref(self) -> queue_pb2.LeaseRef:
+        return queue_pb2.LeaseRef(
+            task_id=self.task.id,
+            attempt=self.attempt,
+            lease_token=self._token,
+        )
+
+    async def _terminal(self, name: str, request: Any) -> None:
+        request.lease.CopyFrom(self._ref())
+        conn = self._queue._client._require_conn()
+        request.context.CopyFrom(conn.request_context())
+        await conn.call(name, request)
+
+    async def ack(self) -> None:
+        """Acknowledge success for the current attempt/lease token."""
+        await self._terminal("Ack", queue_pb2.AckRequest())
+
+    async def nack(self, error: str = "", *, delay: timedelta | None = None) -> None:
+        """Negative-acknowledge; the task is retried per policy."""
+        request = queue_pb2.NackRequest(error=error)
+        if delay is not None:
+            request.retry_delay_override.CopyFrom(_pb.td_to_dur(delay))
+        await self._terminal("Nack", request)
+
+    async def reject(self, error: str = "") -> None:
+        """Permanently fail the task (DEAD) for the current attempt."""
+        await self._terminal("Reject", queue_pb2.RejectRequest(error=error))
+
+    async def extend_lease(self, extension: timedelta) -> None:
+        """Extend the lease to ``server_now + extension`` (absolute reset)."""
+        if extension <= timedelta(0):
+            raise InvalidArgumentError("extension must be positive")
+        request = queue_pb2.ExtendLeaseRequest()
+        request.extension.CopyFrom(_pb.td_to_dur(extension))
+        request.lease.CopyFrom(self._ref())
+        conn = self._queue._client._require_conn()
+        request.context.CopyFrom(conn.request_context())
+        response = await conn.call("ExtendLease", request)
+        self._lease_until = _pb.ts_to_dt(response.lease.until)