durable-worker 0.2.0__tar.gz → 0.4.0__tar.gz

{durable_worker-0.2.0 → durable_worker-0.4.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: durable-worker
-Version: 0.2.0
+Version: 0.4.0
 Summary: Python worker SDK for nestjs-durable — run durable workflow steps in Python.
 Project-URL: Homepage, https://github.com/DavideCarvalho/nestjs-durable
 Project-URL: Repository, https://github.com/DavideCarvalho/nestjs-durable

durable_worker-0.4.0/durable_worker/__init__.py

@@ -0,0 +1,37 @@
+"""durable-worker — Python SDK for running nestjs-durable remote steps.
+
+A worker registers step handlers by name and processes tasks dispatched by the orchestrator.
+The wire protocol (task in, result out) is plain JSON and identical across languages, so the
+same step name implemented here is callable from a TypeScript workflow via ``ctx.call``.
+"""
+
+from .cancellation import Cancelled, CancellationRegistry
+from .redis_runner import redis_url_from_env, run_redis_worker
+from .routing import reply_target
+from .worker import (
+    FatalError,
+    StepContext,
+    Worker,
+    current_step,
+    log,
+    set_process,
+    sub,
+)
+
+__all__ = [
+    "Worker",
+    "FatalError",
+    "StepContext",
+    "Cancelled",
+    "CancellationRegistry",
+    "reply_target",
+    # Context-local step access — record events from anywhere inside a handler without threading ctx.
+    "current_step",
+    "log",
+    "sub",
+    "set_process",
+    # Transport bootstrap helpers.
+    "run_redis_worker",
+    "redis_url_from_env",
+]
+__version__ = "0.4.0"

durable_worker-0.4.0/durable_worker/cancellation.py

@@ -0,0 +1,49 @@
+"""Cooperative cancellation for the Python worker.
+
+The orchestrator broadcasts a ``{"kind": "cancel", "runId": ...}`` control message across every
+instance (Redis pub/sub, same channel the TypeScript control plane uses). A long-running handler
+can't be force-killed, so cancellation is *cooperative*: a runner subscribes to that channel and
+records the cancelled run ids here, and the handler checks ``ctx.cancelled`` (or calls
+``ctx.raise_if_cancelled()``) at safe points to bail out early instead of finishing work whose
+result will be discarded.
+"""
+
+from __future__ import annotations
+
+import threading
+from typing import Any, Dict
+
+
+class Cancelled(Exception):
+    """Raised by ``StepContext.raise_if_cancelled`` when the run has been cancelled."""
+
+    def __init__(self, run_id: str) -> None:
+        super().__init__(f"run {run_id} was cancelled")
+        self.run_id = run_id
+
+
+class CancellationRegistry:
+    """Thread-safe set of cancelled run ids, fed by a control-channel subscription."""
+
+    def __init__(self) -> None:
+        self._cancelled: set[str] = set()
+        self._lock = threading.Lock()
+
+    def cancel(self, run_id: str) -> None:
+        with self._lock:
+            self._cancelled.add(run_id)
+
+    def is_cancelled(self, run_id: str) -> bool:
+        with self._lock:
+            return run_id in self._cancelled
+
+    def clear(self, run_id: str) -> None:
+        with self._lock:
+            self._cancelled.discard(run_id)
+
+    def on_control_message(self, msg: Dict[str, Any]) -> None:
+        """Feed a decoded control message; records the run id when it's a cancel."""
+        if isinstance(msg, dict) and msg.get("kind") == "cancel":
+            run_id = msg.get("runId")
+            if isinstance(run_id, str):
+                self.cancel(run_id)

durable_worker-0.4.0/durable_worker/redis_runner.py

@@ -0,0 +1,173 @@
+"""Run a :class:`Worker` against the BullMQ/Redis transport.
+
+Consumes the orchestrator's per-group tasks queue and publishes results on the shared results
+queue — the same queues a TypeScript ``BullMQTransport`` uses, so steps interoperate across
+languages. Requires the optional ``bullmq`` extra: ``pip install durable-worker[redis]``.
+
+It also subscribes to the orchestrator's control channel (``<prefix>-control``) so a long handler
+can observe cooperative cancellation via ``ctx.cancelled`` — the cross-language half of
+``engine.cancel``.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import json
+import os
+import socket
+from typing import Any, Callable, Dict, Optional
+
+from .cancellation import CancellationRegistry
+from .worker import Worker
+
+# Stable-ish id for the `from` field of control messages this worker publishes. It only has to
+# DIFFER from the engine instanceIds (so a dashboard engine doesn't treat our progress events as its
+# own echo and drop them) — host + pid is plenty and avoids importing a uuid/random dependency.
+_INSTANCE_ID = f"py-{socket.gethostname()}-{os.getpid()}"
+
+
+def redis_url_from_env(prefix: str = "REDIS") -> str:
+    """Build a ``redis://`` URL from ``{prefix}_HOST/_PORT/_USERNAME/_PASSWORD`` env vars. The
+    credentials are URL-encoded — a generated password often contains ``@ : /`` which would corrupt
+    the netloc if left raw. Defaults to ``localhost:6379`` with no auth. Handy for ``Worker.run``."""
+    from urllib.parse import quote
+
+    host = os.getenv(f"{prefix}_HOST", "localhost")
+    port = os.getenv(f"{prefix}_PORT", "6379")
+    user = os.getenv(f"{prefix}_USERNAME") or ""
+    password = os.getenv(f"{prefix}_PASSWORD") or ""
+    auth = f"{quote(user, safe='')}:{quote(password, safe='')}@" if (user or password) else ""
+    return f"redis://{auth}{host}:{port}"
+
+
+def _names(prefix: str, group: str) -> tuple[str, str]:
+    # Must match the TS BullMQTransport: '<prefix>-tasks-<group>' and '<prefix>-results'.
+    return f"{prefix}-tasks-{group}", f"{prefix}-results"
+
+
+def _control_channel(prefix: str) -> str:
+    # Mirrors BullMQTransport.controlChannel(): '<prefix>-control'.
+    return f"{prefix}-control"
+
+
+def _progress_message(task: Dict[str, Any], event: Dict[str, Any]) -> str:
+    """A control-plane `{kind:'event'}` carrying a `step.progress` EngineEvent — the same envelope a
+    TS engine publishes, so a dashboard engine re-delivers it to its subscribers (live-tail). Carries
+    the single just-emitted step event; the `at` mirrors EngineEvent's (ms; the TS side `new Date()`s it)."""
+    return json.dumps(
+        {
+            "kind": "event",
+            "from": _INSTANCE_ID,
+            "event": {
+                "type": "step.progress",
+                "runId": task.get("runId"),
+                "seq": task.get("seq"),
+                "name": task.get("name"),
+                "event": event,
+                "at": event.get("at"),
+            },
+        }
+    )
+
+
+async def run_redis_worker(
+    worker: Worker,
+    *,
+    group: str,
+    connection: str = "redis://localhost:6379",
+    prefix: str = "durable",
+    cancellation: Optional[CancellationRegistry] = None,
+) -> Any:
+    """Start a BullMQ worker that runs ``worker``'s handlers. Returns the bullmq Worker.
+
+    The returned worker runs in the background; ``await worker.close()`` to stop it. When a
+    :class:`CancellationRegistry` is given (one is created otherwise), the runner subscribes to the
+    control channel and feeds it, so handlers see ``ctx.cancelled``.
+    """
+
+    from bullmq import Queue as BullQueue  # imported lazily so the SDK works without bullmq
+    from bullmq import Worker as BullWorker
+
+    tasks_name, results_name = _names(prefix, group)
+    results = BullQueue(results_name, {"connection": connection})
+    registry = cancellation or CancellationRegistry()
+    await _subscribe_control(connection, prefix, registry)
+    publish_progress = await _progress_publisher(connection, prefix)
+
+    async def process(job: Any, _token: str) -> None:
+        on_event = _make_on_event(job.data, publish_progress) if publish_progress else None
+        result = await worker.aprocess_task(
+            job.data, is_cancelled=registry.is_cancelled, on_event=on_event
+        )
+        await results.add("result", result, {"removeOnComplete": True, "removeOnFail": True})
+
+    return BullWorker(tasks_name, process, {"connection": connection})
+
+
+async def _progress_publisher(
+    connection: str, prefix: str
+) -> Optional[Callable[[str], None]]:
+    """Build a thread-safe `publish(message)` that PUBLISHes on the control channel from the running
+    loop. Returns None if redis pub/sub isn't available (then `step.progress` streaming is simply off
+    — the events still ride back on the final result). The returned callable is safe to call from a
+    handler's executor thread: it hops back onto the loop via ``call_soon_threadsafe`` (the aioredis
+    client is bound to this loop, so publishing must happen there, not from the worker thread)."""
+    try:
+        import redis.asyncio as aioredis  # lazy: only needed when streaming progress
+    except ImportError:
+        return None
+
+    client = aioredis.from_url(connection)
+    channel = _control_channel(prefix)
+    loop = asyncio.get_running_loop()
+
+    def publish(message: str) -> None:
+        async def _send() -> None:
+            try:
+                await client.publish(channel, message)
+            except Exception:  # noqa: BLE001 — live-tail is best-effort; never fail the step
+                pass
+
+        try:
+            loop.call_soon_threadsafe(lambda: loop.create_task(_send()))
+        except RuntimeError:
+            pass  # loop is closing/closed — drop the live event (the result still carries it)
+
+    return publish
+
+
+def _make_on_event(
+    task: Dict[str, Any], publish: Callable[[str], None]
+) -> Callable[[Dict[str, Any]], None]:
+    """Per-task sink: turn each step event into a `step.progress` control message and publish it."""
+
+    def on_event(event: Dict[str, Any]) -> None:
+        publish(_progress_message(task, event))
+
+    return on_event
+
+
+async def _subscribe_control(
+    connection: str, prefix: str, registry: CancellationRegistry
+) -> None:
+    """Best-effort: subscribe to the control channel and feed cancellations into ``registry``.
+    No-op (logged) if redis pub/sub isn't available — cancellation just won't be observed."""
+    try:
+        import redis.asyncio as aioredis  # lazy: only needed for cooperative cancellation
+    except ImportError:
+        return
+
+    client = aioredis.from_url(connection)
+    pubsub = client.pubsub()
+    await pubsub.subscribe(_control_channel(prefix))
+
+    async def listen() -> None:
+        async for message in pubsub.listen():
+            if message.get("type") != "message":
+                continue
+            try:
+                registry.on_control_message(json.loads(message["data"]))
+            except (ValueError, TypeError):
+                pass  # ignore malformed control messages
+
+    asyncio.create_task(listen())

durable_worker-0.4.0/durable_worker/routing.py

@@ -0,0 +1,27 @@
+"""Reply routing for a worker that consumes several transports (multi-broker failover).
+
+The orchestrator stamps each task with the id of the transport it was dispatched on
+(``task["transport"]``). A worker that consumes more than one broker must send the *result* back on
+the same broker — so failover is symmetric and the worker never picks a transport itself, it just
+replies where it was told. This is a pure lookup over a ``{id: reply_target}`` mapping with a sane
+fallback for the single-transport / unknown-id case.
+"""
+
+from __future__ import annotations
+
+from typing import Any, Dict, Optional, TypeVar
+
+T = TypeVar("T")
+
+
+def reply_target(task: Dict[str, Any], targets: Dict[str, T]) -> Optional[T]:
+    """The reply target for ``task`` by its ``transport`` id; falls back to the lone target when the
+    id is absent (single transport) or unknown. Returns ``None`` only if ``targets`` is empty."""
+    if not targets:
+        return None
+    transport_id = task.get("transport")
+    if isinstance(transport_id, str) and transport_id in targets:
+        return targets[transport_id]
+    if len(targets) == 1:
+        return next(iter(targets.values()))
+    return None

durable_worker-0.4.0/durable_worker/worker.py

@@ -0,0 +1,388 @@
+"""Core worker: a name->handler registry and the pure task->result dispatch.
+
+Transport (Redis/BullMQ/NATS) is intentionally separate — `process_task` is a pure function of
+the task, so it is fully testable without any broker. A transport adapter just feeds tasks in
+and ships results out.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import contextvars
+import inspect
+import time
+from typing import Any, Awaitable, Callable, Dict, List, Optional, Union
+
+Handler = Callable[..., Union[Any, Awaitable[Any]]]
+
+# The step a handler is currently running, bound for the duration of the call (in the loop for an
+# async handler, in the executor thread for a blocking one). Lets code DEEP inside a handler — that
+# never received `ctx` — record events via the module-level helpers below, without threading the
+# context through every call. Outside a step it's None and the helpers are no-ops.
+_current_step: "contextvars.ContextVar[Optional[StepContext]]" = contextvars.ContextVar(
+    "durable_current_step", default=None
+)
+
+
+def current_step() -> "Optional[StepContext]":
+    """The :class:`StepContext` of the step running on this task/thread, or None outside a step."""
+    return _current_step.get()
+
+
+def log(level: str, message: str, data: Any = None) -> None:
+    """Record a log line on the current step (level: debug/info/warn/error). No-op outside a step."""
+    ctx = _current_step.get()
+    if ctx is not None:
+        getattr(ctx, level, ctx.info)(message, data)
+
+
+def sub(name: str, status: str, message: Optional[str] = None, data: Any = None) -> None:
+    """Record a sub-process outcome (ok/failed/skipped) on the current step. No-op outside a step."""
+    ctx = _current_step.get()
+    if ctx is not None:
+        ctx.sub(name, status, message, data)
+
+
+def set_process(name: Optional[str]) -> None:
+    """Tag the current step's subsequent log lines with sub-process ``name`` (None clears). No-op
+    outside a step — so the same business code runs on a non-durable path untouched."""
+    ctx = _current_step.get()
+    if ctx is not None:
+        ctx.process(name)
+
+
+class FatalError(Exception):
+    """Raise inside a handler to signal a non-retryable failure (mirrors the TS ``FatalError``).
+
+    The engine will not retry the step regardless of its ``retries`` setting.
+    """
+
+    def __init__(self, message: str, code: str | None = None) -> None:
+        super().__init__(message)
+        self.code = code
+
+
+class StepContext:
+    """Lets a handler record what happened inside a step — debug/info/warn/error lines and
+    per-sub-process outcomes. The events ride back on the result as ``events`` and the engine
+    checkpoints them, so the dashboard shows them under the step. The TypeScript counterpart is
+    the ``StepLogger`` handed to ``ctx.step`` — same ``StepEvent`` shape, so observability is
+    symmetric regardless of which language ran the step.
+
+    A handler opts in by declaring a second parameter::
+
+        @worker.step("processing")
+        def run(data, ctx):
+            for proc in data["procs"]:
+                ok = run_proc(proc)
+                ctx.sub(proc["name"], "ok" if ok else "failed")
+            return {"context": {...}}
+    """
+
+    def __init__(
+        self,
+        run_id: Optional[str] = None,
+        is_cancelled: Optional[Callable[[str], bool]] = None,
+        seq: Optional[int] = None,
+        on_event: Optional[Callable[[Dict[str, Any]], None]] = None,
+    ) -> None:
+        self.events: List[Dict[str, Any]] = []
+        self._run_id = run_id
+        self._is_cancelled = is_cancelled
+        # Step seq + a sink for live progress: when a transport supplies ``on_event``, every event is
+        # ALSO handed to it as it happens (e.g. published on the control plane as a ``step.progress``),
+        # so a dashboard tails a long step line-by-line instead of waiting for the final result.
+        self._seq = seq
+        self._on_event = on_event
+        # The sub-process a handler is currently inside (set via ``process``), stamped onto log lines
+        # so the dashboard groups a fan-out step's trail per sub-process.
+        self._process: Optional[str] = None
+
+    def process(self, name: Optional[str]) -> None:
+        """Mark the sub-process whose work is now running; subsequent log lines are tagged with it
+        (until changed or cleared with ``None``). A no-op for ``sub()`` outcome rows, which name
+        themselves. Lets a fan-out handler attribute its log trail to each sub-process."""
+        self._process = name
+
+    @property
+    def cancelled(self) -> bool:
+        """True once the run has been cancelled (a runner subscribed to the control channel marks
+        it). Check this at safe points in a long handler and return/raise to bail out early."""
+        return bool(
+            self._is_cancelled is not None
+            and self._run_id is not None
+            and self._is_cancelled(self._run_id)
+        )
+
+    def raise_if_cancelled(self) -> None:
+        """Raise :class:`Cancelled` if the run has been cancelled — a one-liner abort for handlers."""
+        if self.cancelled:
+            from .cancellation import Cancelled
+
+            raise Cancelled(self._run_id or "")
+
+    def _emit(
+        self,
+        level: str,
+        message: str,
+        name: Optional[str] = None,
+        status: Optional[str] = None,
+        data: Any = None,
+    ) -> None:
+        event: Dict[str, Any] = {"at": int(time.time() * 1000), "level": level, "message": message}
+        if name is not None:
+            event["name"] = name
+        if status is not None:
+            event["status"] = status
+        # Tag a log line (no outcome status) with the sub-process it was emitted inside, if any.
+        elif self._process is not None:
+            event["process"] = self._process
+        if data is not None:
+            event["data"] = data
+        self.events.append(event)
+        # Live progress: hand the event to the transport's sink so it streams now (best-effort — a
+        # broken sink must never fail the handler; the event is already captured for the result).
+        if self._on_event is not None:
+            try:
+                self._on_event(event)
+            except Exception:  # noqa: BLE001 — live-tail is best-effort observability
+                pass
+
+    def debug(self, message: str, data: Any = None) -> None:
+        self._emit("debug", message, data=data)
+
+    def info(self, message: str, data: Any = None) -> None:
+        self._emit("info", message, data=data)
+
+    def warn(self, message: str, data: Any = None) -> None:
+        self._emit("warn", message, data=data)
+
+    def error(self, message: str, data: Any = None) -> None:
+        self._emit("error", message, data=data)
+
+    def sub(
+        self, name: str, status: str, message: Optional[str] = None, data: Any = None
+    ) -> None:
+        """Record a sub-step / sub-process outcome (e.g. one of N parallel p-processes)."""
+        level = "error" if status == "failed" else "warn" if status == "skipped" else "info"
+        self._emit(level, message or name, name=name, status=status, data=data)
+
+
+class Worker:
+    """Registers step handlers by name and turns a dispatched task into a result.
+
+    Example::
+
+        worker = Worker(group="payments")
+
+        @worker.step("payments.charge-card")
+        async def charge(data):
+            res = await stripe.charge(data["orderId"], data["amountCents"])
+            return {"chargeId": res.id}
+    """
+
+    def __init__(self, group: str = "default") -> None:
+        self.group = group
+        self._handlers: Dict[str, Handler] = {}
+        self._blocking: Dict[str, bool] = {}
+
+    def step(self, name: str, *, blocking: bool = False) -> Callable[[Handler], Handler]:
+        """Decorator registering ``fn`` as the handler for step ``name``.
+
+        Set ``blocking=True`` for a synchronous handler that does CPU/DB work: the worker runs it in
+        a thread so the event loop stays free to renew the broker's job lock (otherwise a long step
+        looks stalled and gets redelivered). The step context is bound inside that thread too, so the
+        module-level :func:`log`/:func:`sub`/:func:`set_process` work from anywhere in the call.
+        """
+
+        def register(fn: Handler) -> Handler:
+            self._handlers[name] = fn
+            self._blocking[name] = blocking
+            return fn
+
+        return register
+
+    def handles(self, name: str) -> bool:
+        return name in self._handlers
+
+    def run(self, *, redis: str = "redis://localhost:6379", prefix: str = "durable") -> None:
+        """Run this worker against the Redis/BullMQ transport until the process is signalled, then
+        close gracefully. Owns the whole long-running bootstrap a worker process needs — the event
+        loop, SIGTERM/SIGINT handling, and the broker connection — so the entrypoint is one call::
+
+            worker = Worker(group="processing")
+
+            @worker.step("processing", blocking=True)
+            def handle(data, ctx): ...
+
+            worker.run(redis=redis_url_from_env())
+
+        Blocks until a shutdown signal closes the worker (which finishes the active job and releases
+        its lock). For finer control (your own loop/signals), use :func:`run_redis_worker` directly.
+        """
+        import signal as _signal
+
+        from .redis_runner import run_redis_worker
+
+        async def _main() -> None:
+            bull_worker = await run_redis_worker(
+                self, group=self.group, connection=redis, prefix=prefix
+            )
+            stop = asyncio.Event()
+            loop = asyncio.get_running_loop()
+
+            async def _graceful() -> None:
+                try:
+                    await bull_worker.close()
+                finally:
+                    stop.set()
+
+            for sig in (_signal.SIGTERM, _signal.SIGINT):
+                try:
+                    loop.add_signal_handler(sig, lambda: asyncio.ensure_future(_graceful()))
+                except (NotImplementedError, RuntimeError):
+                    pass  # no signal support here (e.g. not the main thread) — rely on redelivery
+            await stop.wait()
+
+        asyncio.run(_main())
+
+    def process_task(
+        self,
+        task: Dict[str, Any],
+        is_cancelled: Optional[Callable[[str], bool]] = None,
+        on_event: Optional[Callable[[Dict[str, Any]], None]] = None,
+    ) -> Dict[str, Any]:
+        """Run the handler for ``task`` and return a wire-format result.
+
+        Pure and synchronous from the caller's view (async handlers are awaited internally), so a
+        transport can simply ``result = worker.process_task(task); send(result)``. Pass
+        ``is_cancelled`` so the handler's ``ctx.cancelled`` reflects cooperative cancellation, and
+        ``on_event`` to receive each step event live (for ``step.progress`` streaming).
+        """
+
+        base = self._base(task)
+        handler = self._handlers.get(task["name"])
+        if handler is None:
+            return self._no_handler(base, task["name"])
+        ctx = self._ctx(task, is_cancelled, on_event)
+        try:
+            output = _run_bound(handler, task.get("input"), ctx)
+            if inspect.isawaitable(output):
+                output = asyncio.run(output)
+            return self._completed(base, output, ctx)
+        except Exception as err:  # noqa: BLE001
+            return self._failure(base, err, ctx)
+
+    async def aprocess_task(
+        self,
+        task: Dict[str, Any],
+        is_cancelled: Optional[Callable[[str], bool]] = None,
+        on_event: Optional[Callable[[Dict[str, Any]], None]] = None,
+    ) -> Dict[str, Any]:
+        """Async variant — awaits async handlers in the current loop. Use from a transport that
+        already runs inside an event loop (e.g. the BullMQ runner). A handler registered
+        ``blocking=True`` runs in a thread so the loop stays free to renew the job lock."""
+
+        base = self._base(task)
+        handler = self._handlers.get(task["name"])
+        if handler is None:
+            return self._no_handler(base, task["name"])
+        ctx = self._ctx(task, is_cancelled, on_event)
+        try:
+            if self._blocking.get(task["name"]):
+                loop = asyncio.get_running_loop()
+                output = await loop.run_in_executor(
+                    None, _run_bound, handler, task.get("input"), ctx
+                )
+            else:
+                output = _run_bound(handler, task.get("input"), ctx)
+                if inspect.isawaitable(output):
+                    # The context var was reset by _run_bound after the sync part returned the
+                    # coroutine; re-bind it for the duration of the await so deep `log()` calls work.
+                    token = _current_step.set(ctx)
+                    try:
+                        output = await output
+                    finally:
+                        _current_step.reset(token)
+            return self._completed(base, output, ctx)
+        except Exception as err:  # noqa: BLE001
+            return self._failure(base, err, ctx)
+
+    @staticmethod
+    def _ctx(
+        task: Dict[str, Any],
+        is_cancelled: Optional[Callable[[str], bool]],
+        on_event: Optional[Callable[[Dict[str, Any]], None]],
+    ) -> StepContext:
+        return StepContext(
+            run_id=task.get("runId"),
+            is_cancelled=is_cancelled,
+            seq=task.get("seq"),
+            on_event=on_event,
+        )
+
+    @staticmethod
+    def _base(task: Dict[str, Any]) -> Dict[str, Any]:
+        # Stamp the worker's pickup time so the engine can report queue-wait (startedAt − enqueuedAt),
+        # mirroring the TypeScript ``runStepHandler``. Epoch ms keeps it language-neutral on the wire.
+        return {
+            "runId": task["runId"],
+            "seq": task["seq"],
+            "stepId": task["stepId"],
+            "startedAt": int(time.time() * 1000),
+        }
+
+    @staticmethod
+    def _completed(base: Dict[str, Any], output: Any, ctx: StepContext) -> Dict[str, Any]:
+        return _with_events({**base, "status": "completed", "output": output}, ctx)
+
+    @staticmethod
+    def _no_handler(base: Dict[str, Any], name: str) -> Dict[str, Any]:
+        return {
+            **base,
+            "status": "failed",
+            "error": {"message": f"no handler for {name}", "retryable": False},
+        }
+
+    @staticmethod
+    def _failure(base: Dict[str, Any], err: Exception, ctx: StepContext) -> Dict[str, Any]:
+        # Keep whatever the handler logged before it threw — a partial p-process run is exactly
+        # the case where the sub-process outcomes matter most.
+        if isinstance(err, FatalError):
+            error = {"message": str(err), "code": err.code, "retryable": False}
+        else:
+            error = {"message": str(err)}
+        return _with_events({**base, "status": "failed", "error": error}, ctx)
+
+
+def _run_bound(handler: Handler, input_: Any, ctx: "StepContext") -> Any:
+    """Call ``handler`` with ``ctx`` bound as the current step for the duration of the call, so the
+    module-level :func:`log`/:func:`sub`/:func:`set_process` reach it from deep in the call tree.
+    Passes ``ctx`` as a second arg only if the handler declares one (plain ``def h(data)`` still
+    works). Also the target run in the executor thread for a ``blocking=True`` handler."""
+    token = _current_step.set(ctx)
+    try:
+        return handler(input_, ctx) if _wants_ctx(handler) else handler(input_)
+    finally:
+        _current_step.reset(token)
+
+
+def _wants_ctx(handler: Handler) -> bool:
+    """True if ``handler`` can accept the step context as a second positional argument."""
+    try:
+        params = list(inspect.signature(handler).parameters.values())
+    except (ValueError, TypeError):  # builtins / C functions without a signature
+        return False
+    positional = [
+        p
+        for p in params
+        if p.kind in (inspect.Parameter.POSITIONAL_ONLY, inspect.Parameter.POSITIONAL_OR_KEYWORD)
+    ]
+    has_var_positional = any(p.kind is inspect.Parameter.VAR_POSITIONAL for p in params)
+    return len(positional) >= 2 or has_var_positional
+
+
+def _with_events(result: Dict[str, Any], ctx: StepContext) -> Dict[str, Any]:
+    if ctx.events:
+        result["events"] = ctx.events
+    return result

{durable_worker-0.2.0 → durable_worker-0.4.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "durable-worker"
-version = "0.2.0"
+version = "0.4.0"
 description = "Python worker SDK for nestjs-durable — run durable workflow steps in Python."
 readme = "README.md"
 requires-python = ">=3.9"

durable_worker-0.4.0/tests/test_cancellation.py

@@ -0,0 +1,68 @@
+import unittest
+
+from durable_worker import CancellationRegistry, Cancelled, Worker
+
+
+def task(name="long.job", **over):
+    base = {"runId": "r1", "seq": 0, "name": name, "stepId": "r1:0", "input": {}, "attempt": 1}
+    base.update(over)
+    return base
+
+
+class CancellationRegistryTest(unittest.TestCase):
+    def test_records_a_cancel_control_message(self):
+        reg = CancellationRegistry()
+        self.assertFalse(reg.is_cancelled("r1"))
+        reg.on_control_message({"kind": "cancel", "runId": "r1", "from": "x"})
+        self.assertTrue(reg.is_cancelled("r1"))
+
+    def test_ignores_non_cancel_messages(self):
+        reg = CancellationRegistry()
+        reg.on_control_message({"kind": "event", "event": {}})
+        reg.on_control_message({"nonsense": True})
+        self.assertFalse(reg.is_cancelled("r1"))
+
+
+class StepContextCancellationTest(unittest.TestCase):
+    def test_handler_sees_cancellation_via_ctx(self):
+        reg = CancellationRegistry()
+        worker = Worker()
+        observed = {}
+
+        @worker.step("long.job")
+        def run(_data, ctx):
+            observed["before"] = ctx.cancelled
+            reg.cancel("r1")
+            observed["after"] = ctx.cancelled
+            return "done"
+
+        worker.process_task(task(), is_cancelled=reg.is_cancelled)
+        self.assertEqual(observed, {"before": False, "after": True})
+
+    def test_raise_if_cancelled_aborts_the_step(self):
+        reg = CancellationRegistry()
+        reg.cancel("r1")
+        worker = Worker()
+
+        @worker.step("long.job")
+        def run(_data, ctx):
+            ctx.raise_if_cancelled()
+            return "should not get here"
+
+        result = worker.process_task(task(), is_cancelled=reg.is_cancelled)
+        self.assertEqual(result["status"], "failed")
+        self.assertIn("cancelled", result["error"]["message"])
+
+    def test_cancelled_is_false_without_a_source(self):
+        worker = Worker()
+
+        @worker.step("long.job")
+        def run(_data, ctx):
+            return ctx.cancelled
+
+        result = worker.process_task(task())
+        self.assertEqual(result["output"], False)
+
+
+if __name__ == "__main__":
+    unittest.main()

durable_worker-0.4.0/tests/test_redis_url.py

@@ -0,0 +1,51 @@
+import unittest
+
+from durable_worker import redis_url_from_env
+
+
+class RedisUrlFromEnvTest(unittest.TestCase):
+    def setUp(self):
+        import os
+
+        self._saved = {
+            k: os.environ.pop(k, None)
+            for k in ("REDIS_HOST", "REDIS_PORT", "REDIS_USERNAME", "REDIS_PASSWORD")
+        }
+
+    def tearDown(self):
+        import os
+
+        for k, v in self._saved.items():
+            if v is None:
+                os.environ.pop(k, None)
+            else:
+                os.environ[k] = v
+
+    def test_defaults_to_localhost_without_auth(self):
+        self.assertEqual(redis_url_from_env(), "redis://localhost:6379")
+
+    def test_host_port_and_auth(self):
+        import os
+
+        os.environ["REDIS_HOST"] = "redis.internal"
+        os.environ["REDIS_PORT"] = "6380"
+        self.assertEqual(redis_url_from_env(), "redis://redis.internal:6380")
+        os.environ["REDIS_USERNAME"] = "default"
+        os.environ["REDIS_PASSWORD"] = "secret"
+        self.assertEqual(redis_url_from_env(), "redis://default:secret@redis.internal:6380")
+
+    def test_url_encodes_credentials(self):
+        import os
+
+        os.environ["REDIS_HOST"] = "redis.internal"
+        os.environ["REDIS_PORT"] = "6380"
+        os.environ["REDIS_USERNAME"] = "default"
+        # A password with @ : / would corrupt the netloc if left raw.
+        os.environ["REDIS_PASSWORD"] = "p@ss:w/rd"
+        self.assertEqual(
+            redis_url_from_env(), "redis://default:p%40ss%3Aw%2Frd@redis.internal:6380"
+        )
+
+
+if __name__ == "__main__":
+    unittest.main()

durable_worker-0.4.0/tests/test_routing.py

@@ -0,0 +1,23 @@
+import unittest
+
+from durable_worker import reply_target
+
+
+class ReplyTargetTest(unittest.TestCase):
+    def test_routes_by_task_transport_id(self):
+        targets = {"redis": "redis-results", "sqs": "sqs-results"}
+        self.assertEqual(reply_target({"transport": "sqs"}, targets), "sqs-results")
+        self.assertEqual(reply_target({"transport": "redis"}, targets), "redis-results")
+
+    def test_falls_back_to_the_lone_target_when_id_is_absent(self):
+        self.assertEqual(reply_target({}, {"only": "x"}), "x")
+
+    def test_unknown_id_with_multiple_targets_is_unroutable(self):
+        self.assertIsNone(reply_target({"transport": "nats"}, {"redis": "a", "sqs": "b"}))
+
+    def test_empty_targets_is_none(self):
+        self.assertIsNone(reply_target({"transport": "redis"}, {}))
+
+
+if __name__ == "__main__":
+    unittest.main()

{durable_worker-0.2.0 → durable_worker-0.4.0}/tests/test_worker.py

@@ -1,6 +1,15 @@
+import asyncio
 import unittest
 
-from durable_worker import FatalError, StepContext, Worker
+from durable_worker import (
+    FatalError,
+    StepContext,
+    Worker,
+    current_step,
+    log,
+    set_process,
+    sub,
+)
 
 
 def task(name="payments.charge-card", **over):
@@ -136,6 +145,97 @@ class WorkerTest(unittest.TestCase):
         result = worker.process_task(task())
         self.assertEqual(result["events"][0]["message"], "async")
 
+    def test_logs_are_tagged_with_the_current_process(self):
+        worker = Worker()
+
+        @worker.step("payments.charge-card")
+        def charge(_data, ctx: StepContext):
+            ctx.process("proc-a")
+            ctx.info("inside a")
+            ctx.sub("proc-a", "ok")
+            ctx.process("proc-b")
+            ctx.debug("inside b")
+            ctx.process(None)
+            ctx.info("step-level")
+            return {"ok": True}
+
+        # process() itself emits nothing, so the events are: [info a, sub a, debug b, info step-level]
+        events = worker.process_task(task())["events"]
+        # log lines carry the sub-process that was running when they fired...
+        self.assertEqual(events[0]["process"], "proc-a")
+        self.assertEqual(events[2]["process"], "proc-b")
+        # ...the sub-process OUTCOME row names itself (no `process` tag)...
+        self.assertNotIn("process", events[1])
+        self.assertEqual(events[1]["status"], "ok")
+        # ...and a log after process(None) is untagged (step-level).
+        self.assertNotIn("process", events[3])
+
+    def test_on_event_streams_each_event_as_it_is_emitted(self):
+        worker = Worker()
+        live = []
+
+        @worker.step("payments.charge-card")
+        def charge(_data, ctx: StepContext):
+            ctx.info("first")
+            # the live sink saw the event the instant it was emitted, before the step returned
+            self.assertEqual([e["message"] for e in live], ["first"])
+            ctx.sub("proc-a", "ok")
+            return {"ok": True}
+
+        result = worker.process_task(task(), on_event=live.append)
+        self.assertEqual([e["message"] for e in live], ["first", "proc-a"])
+        # the same events still ride back on the result (live-tail doesn't replace the final record)
+        self.assertEqual(len(result["events"]), 2)
+
+
+    def test_module_level_helpers_reach_the_current_step(self):
+        worker = Worker()
+
+        def deep_business_logic():
+            # No `ctx` in scope here — the context-local helpers find it.
+            set_process("proc-a")
+            log("debug", "querying")
+            sub("proc-a", "ok")
+
+        @worker.step("payments.charge-card")
+        def charge(_data):  # note: no ctx param — handler doesn't even take it
+            deep_business_logic()
+            return {"ok": True}
+
+        events = worker.process_task(task())["events"]
+        self.assertEqual(events[0]["message"], "querying")
+        self.assertEqual(events[0]["process"], "proc-a")
+        self.assertEqual(events[1]["name"], "proc-a")
+        self.assertEqual(events[1]["status"], "ok")
+
+    def test_module_level_helpers_are_noops_outside_a_step(self):
+        # Same business code on a non-durable path must not blow up.
+        self.assertIsNone(current_step())
+        log("info", "no step here")
+        sub("x", "ok")
+        set_process("x")
+
+    def test_blocking_handler_runs_in_a_thread_with_the_step_bound(self):
+        import threading
+
+        worker = Worker()
+        ran_on = {}
+
+        @worker.step("payments.charge-card", blocking=True)
+        def charge(_data):
+            ran_on["thread"] = threading.current_thread().name
+            log("info", "from the worker thread")
+            return {"ok": True}
+
+        async def drive():
+            return await worker.aprocess_task(task())
+
+        loop_thread = threading.current_thread().name
+        result = asyncio.run(drive())
+        self.assertEqual(result["events"][0]["message"], "from the worker thread")
+        # it really ran off the event-loop thread (so the lock keeps renewing)
+        self.assertNotEqual(ran_on["thread"], loop_thread)
+
 
 if __name__ == "__main__":
     unittest.main()

durable_worker-0.2.0/durable_worker/__init__.py

@@ -1,11 +0,0 @@
-"""durable-worker — Python SDK for running nestjs-durable remote steps.
-
-A worker registers step handlers by name and processes tasks dispatched by the orchestrator.
-The wire protocol (task in, result out) is plain JSON and identical across languages, so the
-same step name implemented here is callable from a TypeScript workflow via ``ctx.call``.
-"""
-
-from .worker import FatalError, StepContext, Worker
-
-__all__ = ["Worker", "FatalError", "StepContext"]
-__version__ = "0.2.0"

durable_worker-0.2.0/durable_worker/redis_runner.py

@@ -1,42 +0,0 @@
-"""Run a :class:`Worker` against the BullMQ/Redis transport.
-
-Consumes the orchestrator's per-group tasks queue and publishes results on the shared results
-queue — the same queues a TypeScript ``BullMQTransport`` uses, so steps interoperate across
-languages. Requires the optional ``bullmq`` extra: ``pip install durable-worker[redis]``.
-"""
-
-from __future__ import annotations
-
-from typing import Any
-
-from .worker import Worker
-
-
-def _names(prefix: str, group: str) -> tuple[str, str]:
-    # Must match the TS BullMQTransport: '<prefix>-tasks-<group>' and '<prefix>-results'.
-    return f"{prefix}-tasks-{group}", f"{prefix}-results"
-
-
-async def run_redis_worker(
-    worker: Worker,
-    *,
-    group: str,
-    connection: str = "redis://localhost:6379",
-    prefix: str = "durable",
-) -> Any:
-    """Start a BullMQ worker that runs ``worker``'s handlers. Returns the bullmq Worker.
-
-    The returned worker runs in the background; ``await worker.close()`` to stop it.
-    """
-
-    from bullmq import Queue as BullQueue  # imported lazily so the SDK works without bullmq
-    from bullmq import Worker as BullWorker
-
-    tasks_name, results_name = _names(prefix, group)
-    results = BullQueue(results_name, {"connection": connection})
-
-    async def process(job: Any, _token: str) -> None:
-        result = await worker.aprocess_task(job.data)
-        await results.add("result", result, {"removeOnComplete": True, "removeOnFail": True})
-
-    return BullWorker(tasks_name, process, {"connection": connection})

durable_worker-0.2.0/durable_worker/worker.py

@@ -1,212 +0,0 @@
-"""Core worker: a name->handler registry and the pure task->result dispatch.
-
-Transport (Redis/BullMQ/NATS) is intentionally separate — `process_task` is a pure function of
-the task, so it is fully testable without any broker. A transport adapter just feeds tasks in
-and ships results out.
-"""
-
-from __future__ import annotations
-
-import asyncio
-import inspect
-import time
-from typing import Any, Awaitable, Callable, Dict, List, Optional, Union
-
-Handler = Callable[..., Union[Any, Awaitable[Any]]]
-
-
-class FatalError(Exception):
-    """Raise inside a handler to signal a non-retryable failure (mirrors the TS ``FatalError``).
-
-    The engine will not retry the step regardless of its ``retries`` setting.
-    """
-
-    def __init__(self, message: str, code: str | None = None) -> None:
-        super().__init__(message)
-        self.code = code
-
-
-class StepContext:
-    """Lets a handler record what happened inside a step — debug/info/warn/error lines and
-    per-sub-process outcomes. The events ride back on the result as ``events`` and the engine
-    checkpoints them, so the dashboard shows them under the step. The TypeScript counterpart is
-    the ``StepLogger`` handed to ``ctx.step`` — same ``StepEvent`` shape, so observability is
-    symmetric regardless of which language ran the step.
-
-    A handler opts in by declaring a second parameter::
-
-        @worker.step("processing")
-        def run(data, ctx):
-            for proc in data["procs"]:
-                ok = run_proc(proc)
-                ctx.sub(proc["name"], "ok" if ok else "failed")
-            return {"context": {...}}
-    """
-
-    def __init__(self) -> None:
-        self.events: List[Dict[str, Any]] = []
-
-    def _emit(
-        self,
-        level: str,
-        message: str,
-        name: Optional[str] = None,
-        status: Optional[str] = None,
-        data: Any = None,
-    ) -> None:
-        event: Dict[str, Any] = {"at": int(time.time() * 1000), "level": level, "message": message}
-        if name is not None:
-            event["name"] = name
-        if status is not None:
-            event["status"] = status
-        if data is not None:
-            event["data"] = data
-        self.events.append(event)
-
-    def debug(self, message: str, data: Any = None) -> None:
-        self._emit("debug", message, data=data)
-
-    def info(self, message: str, data: Any = None) -> None:
-        self._emit("info", message, data=data)
-
-    def warn(self, message: str, data: Any = None) -> None:
-        self._emit("warn", message, data=data)
-
-    def error(self, message: str, data: Any = None) -> None:
-        self._emit("error", message, data=data)
-
-    def sub(
-        self, name: str, status: str, message: Optional[str] = None, data: Any = None
-    ) -> None:
-        """Record a sub-step / sub-process outcome (e.g. one of N parallel p-processes)."""
-        level = "error" if status == "failed" else "warn" if status == "skipped" else "info"
-        self._emit(level, message or name, name=name, status=status, data=data)
-
-
-class Worker:
-    """Registers step handlers by name and turns a dispatched task into a result.
-
-    Example::
-
-        worker = Worker(group="payments")
-
-        @worker.step("payments.charge-card")
-        async def charge(data):
-            res = await stripe.charge(data["orderId"], data["amountCents"])
-            return {"chargeId": res.id}
-    """
-
-    def __init__(self, group: str = "default") -> None:
-        self.group = group
-        self._handlers: Dict[str, Handler] = {}
-
-    def step(self, name: str) -> Callable[[Handler], Handler]:
-        """Decorator registering ``fn`` as the handler for step ``name``."""
-
-        def register(fn: Handler) -> Handler:
-            self._handlers[name] = fn
-            return fn
-
-        return register
-
-    def handles(self, name: str) -> bool:
-        return name in self._handlers
-
-    def process_task(self, task: Dict[str, Any]) -> Dict[str, Any]:
-        """Run the handler for ``task`` and return a wire-format result.
-
-        Pure and synchronous from the caller's view (async handlers are awaited internally), so a
-        transport can simply ``result = worker.process_task(task); send(result)``.
-        """
-
-        base = self._base(task)
-        handler = self._handlers.get(task["name"])
-        if handler is None:
-            return self._no_handler(base, task["name"])
-        ctx = StepContext()
-        try:
-            output = self._invoke(handler, task.get("input"), ctx)
-            if inspect.isawaitable(output):
-                output = asyncio.run(output)
-            return self._completed(base, output, ctx)
-        except Exception as err:  # noqa: BLE001
-            return self._failure(base, err, ctx)
-
-    async def aprocess_task(self, task: Dict[str, Any]) -> Dict[str, Any]:
-        """Async variant — awaits async handlers in the current loop. Use from a transport that
-        already runs inside an event loop (e.g. the BullMQ runner)."""
-
-        base = self._base(task)
-        handler = self._handlers.get(task["name"])
-        if handler is None:
-            return self._no_handler(base, task["name"])
-        ctx = StepContext()
-        try:
-            output = self._invoke(handler, task.get("input"), ctx)
-            if inspect.isawaitable(output):
-                output = await output
-            return self._completed(base, output, ctx)
-        except Exception as err:  # noqa: BLE001
-            return self._failure(base, err, ctx)
-
-    @staticmethod
-    def _invoke(handler: Handler, input_: Any, ctx: StepContext) -> Any:
-        """Call ``handler`` with the step input, passing ``ctx`` only if it declares a second
-        parameter — so plain ``def handler(data)`` handlers keep working unchanged."""
-        if _wants_ctx(handler):
-            return handler(input_, ctx)
-        return handler(input_)
-
-    @staticmethod
-    def _base(task: Dict[str, Any]) -> Dict[str, Any]:
-        # Stamp the worker's pickup time so the engine can report queue-wait (startedAt − enqueuedAt),
-        # mirroring the TypeScript ``runStepHandler``. Epoch ms keeps it language-neutral on the wire.
-        return {
-            "runId": task["runId"],
-            "seq": task["seq"],
-            "stepId": task["stepId"],
-            "startedAt": int(time.time() * 1000),
-        }
-
-    @staticmethod
-    def _completed(base: Dict[str, Any], output: Any, ctx: StepContext) -> Dict[str, Any]:
-        return _with_events({**base, "status": "completed", "output": output}, ctx)
-
-    @staticmethod
-    def _no_handler(base: Dict[str, Any], name: str) -> Dict[str, Any]:
-        return {
-            **base,
-            "status": "failed",
-            "error": {"message": f"no handler for {name}", "retryable": False},
-        }
-
-    @staticmethod
-    def _failure(base: Dict[str, Any], err: Exception, ctx: StepContext) -> Dict[str, Any]:
-        # Keep whatever the handler logged before it threw — a partial p-process run is exactly
-        # the case where the sub-process outcomes matter most.
-        if isinstance(err, FatalError):
-            error = {"message": str(err), "code": err.code, "retryable": False}
-        else:
-            error = {"message": str(err)}
-        return _with_events({**base, "status": "failed", "error": error}, ctx)
-
-
-def _wants_ctx(handler: Handler) -> bool:
-    """True if ``handler`` can accept the step context as a second positional argument."""
-    try:
-        params = list(inspect.signature(handler).parameters.values())
-    except (ValueError, TypeError):  # builtins / C functions without a signature
-        return False
-    positional = [
-        p
-        for p in params
-        if p.kind in (inspect.Parameter.POSITIONAL_ONLY, inspect.Parameter.POSITIONAL_OR_KEYWORD)
-    ]
-    has_var_positional = any(p.kind is inspect.Parameter.VAR_POSITIONAL for p in params)
-    return len(positional) >= 2 or has_var_positional
-
-
-def _with_events(result: Dict[str, Any], ctx: StepContext) -> Dict[str, Any]:
-    if ctx.events:
-        result["events"] = ctx.events
-    return result