Add Python worker interceptors

Issue: zorporation/durable-workflow#451

Loop-ID: build-03

Author: durable-workflow-ops

README.md

@@ -141,6 +141,7 @@ advertise support for the target workflow type.
 - **HTTP/JSON protocol**: No gRPC, no protobuf dependencies
 - **Codec envelopes**: Avro payloads by default, with JSON decode compatibility for existing history
 - **Payload-size warnings**: Structured warnings before oversized workflow, activity, signal, update, query, or search-attribute payloads reach the server
+- **Worker interceptors**: Typed hooks around workflow tasks, activity calls, and query tasks for tracing, logging, and custom metrics
 - **Metrics hooks**: Pluggable counters and histograms, with an optional Prometheus adapter
 
 ## Payload-size warnings
@@ -207,6 +208,44 @@ client = Client("http://server:8080", token="dev-token-123", metrics=metrics)
 
 Custom recorders implement `increment(name, value=1.0, tags=None)` and `record(name, value, tags=None)`.
 
+## Worker interceptors
+
+Use `Worker(interceptors=[...])` when instrumentation needs the task payload,
+result, or exception around worker execution instead of only aggregate counters.
+Interceptors run in list order; the first interceptor is the outermost wrapper.
+
+```python
+from durable_workflow import (
+    ActivityInterceptorContext,
+    ActivityHandler,
+    PassthroughWorkerInterceptor,
+    Worker,
+)
+
+class LoggingInterceptor(PassthroughWorkerInterceptor):
+    async def execute_activity(
+        self,
+        context: ActivityInterceptorContext,
+        next: ActivityHandler,
+    ) -> object:
+        print("activity started", context.activity_type)
+        try:
+            result = await next(context)
+        except Exception:
+            print("activity failed", context.activity_type)
+            raise
+        print("activity completed", context.activity_type)
+        return result
+
+worker = Worker(
+    client,
+    task_queue="python-workers",
+    workflows=[GreeterWorkflow],
+    activities=[greet],
+    interceptors=[LoggingInterceptor()],
+)
+```
+
 ## Documentation
 
 Full documentation is available at

src/durable_workflow/__init__.py

@@ -40,6 +40,16 @@
     WorkflowNotFound,
     WorkflowTerminated,
 )
+from .interceptors import (
+    ActivityHandler,
+    ActivityInterceptorContext,
+    PassthroughWorkerInterceptor,
+    QueryTaskHandler,
+    QueryTaskInterceptorContext,
+    WorkerInterceptor,
+    WorkflowTaskHandler,
+    WorkflowTaskInterceptorContext,
+)
 from .metrics import (
     InMemoryMetrics,
     MetricsRecorder,
@@ -55,7 +65,9 @@
     "__version__",
     "ActivityCancelled",
     "ActivityContext",
+    "ActivityHandler",
     "ActivityInfo",
+    "ActivityInterceptorContext",
     "ActivityRetryPolicy",
     "ChildWorkflowRetryPolicy",
     "ChildWorkflowFailed",
@@ -73,7 +85,10 @@
     "ScheduleTriggerResult",
     "StartChildWorkflow",
     "Worker",
+    "WorkerInterceptor",
     "WorkflowExecution",
+    "WorkflowTaskHandler",
+    "WorkflowTaskInterceptorContext",
     "WorkflowHandle",
     "WorkflowList",
     "activity",
@@ -90,6 +105,9 @@
     "PrometheusMetrics",
     "PayloadSizeWarningConfig",
     "PayloadSizeWarningContext",
+    "PassthroughWorkerInterceptor",
+    "QueryTaskHandler",
+    "QueryTaskInterceptorContext",
     "RetryPolicy",
     "ServerError",
     "TransportRetryPolicy",

src/durable_workflow/interceptors.py

@@ -0,0 +1,99 @@
+"""Worker interceptor protocols for SDK instrumentation hooks."""
+
+from __future__ import annotations
+
+from collections.abc import Awaitable, Callable
+from dataclasses import dataclass
+from typing import Any, Protocol
+
+
+@dataclass(frozen=True)
+class WorkflowTaskInterceptorContext:
+    """Context passed to workflow task interceptor hooks."""
+
+    worker_id: str
+    task_queue: str
+    task: dict[str, Any]
+
+
+@dataclass(frozen=True)
+class ActivityInterceptorContext:
+    """Context passed to activity execution interceptor hooks."""
+
+    worker_id: str
+    task_queue: str
+    task: dict[str, Any]
+    activity_type: str
+    args: tuple[Any, ...]
+
+
+@dataclass(frozen=True)
+class QueryTaskInterceptorContext:
+    """Context passed to query task interceptor hooks."""
+
+    worker_id: str
+    task_queue: str
+    task: dict[str, Any]
+
+
+WorkflowTaskHandler = Callable[[WorkflowTaskInterceptorContext], Awaitable[list[dict[str, Any]] | None]]
+ActivityHandler = Callable[[ActivityInterceptorContext], Awaitable[Any]]
+QueryTaskHandler = Callable[[QueryTaskInterceptorContext], Awaitable[str]]
+
+
+class WorkerInterceptor(Protocol):
+    """Protocol for wrapping worker task execution.
+
+    Interceptors run in the order passed to ``Worker(..., interceptors=[...])``.
+    The first interceptor is the outermost wrapper and should call ``next`` to
+    continue the chain.
+    """
+
+    async def execute_workflow_task(
+        self,
+        context: WorkflowTaskInterceptorContext,
+        next: WorkflowTaskHandler,
+    ) -> list[dict[str, Any]] | None:
+        """Wrap workflow task execution."""
+        ...
+
+    async def execute_activity(
+        self,
+        context: ActivityInterceptorContext,
+        next: ActivityHandler,
+    ) -> Any:
+        """Wrap the registered activity callable."""
+        ...
+
+    async def execute_query_task(
+        self,
+        context: QueryTaskInterceptorContext,
+        next: QueryTaskHandler,
+    ) -> str:
+        """Wrap query task execution."""
+        ...
+
+
+class PassthroughWorkerInterceptor:
+    """No-op interceptor useful as a base for partial custom interceptors."""
+
+    async def execute_workflow_task(
+        self,
+        context: WorkflowTaskInterceptorContext,
+        next: WorkflowTaskHandler,
+    ) -> list[dict[str, Any]] | None:
+        return await next(context)
+
+    async def execute_activity(
+        self,
+        context: ActivityInterceptorContext,
+        next: ActivityHandler,
+    ) -> Any:
+        return await next(context)
+
+    async def execute_query_task(
+        self,
+        context: QueryTaskInterceptorContext,
+        next: QueryTaskHandler,
+    ) -> str:
+        return await next(context)

src/durable_workflow/worker.py

@@ -21,7 +21,7 @@
 import time
 import traceback
 import uuid
-from collections.abc import Callable, Iterable
+from collections.abc import Awaitable, Callable, Iterable
 from typing import Any
 
 from . import serializer
@@ -35,6 +35,12 @@
     WorkflowExecution,
 )
 from .errors import ActivityCancelled, AvroNotInstalledError, NonRetryableError, QueryFailed
+from .interceptors import (
+    ActivityInterceptorContext,
+    QueryTaskInterceptorContext,
+    WorkerInterceptor,
+    WorkflowTaskInterceptorContext,
+)
 from .metrics import (
     NOOP_METRICS,
     WORKER_POLL_DURATION_SECONDS,
@@ -151,6 +157,7 @@ def __init__(
         max_concurrent_activity_tasks: int = 10,
         shutdown_timeout: float = 30.0,
         metrics: MetricsRecorder | None = None,
+        interceptors: Iterable[WorkerInterceptor] = (),
     ) -> None:
         self.client = client
         self.task_queue = task_queue
@@ -166,6 +173,7 @@ def __init__(
         self._query_tasks_supported = False
         configured_metrics = metrics if metrics is not None else getattr(client, "metrics", NOOP_METRICS)
         self.metrics: MetricsRecorder = configured_metrics or NOOP_METRICS
+        self.interceptors = tuple(interceptors)
 
     def _record_poll_metrics(self, task_kind: str, outcome: str, duration: float) -> None:
         tags = {"task_kind": task_kind, "task_queue": self.task_queue, "outcome": outcome}
@@ -201,6 +209,35 @@ async def _register(self) -> None:
         log.info("worker %s registered on %s", self.worker_id, self.task_queue)
 
     async def _run_workflow_task(self, task: dict[str, Any]) -> list[dict[str, Any]] | None:
+        context = WorkflowTaskInterceptorContext(
+            worker_id=self.worker_id,
+            task_queue=self.task_queue,
+            task=task,
+        )
+
+        async def call_core(ctx: WorkflowTaskInterceptorContext) -> list[dict[str, Any]] | None:
+            return await self._run_workflow_task_core(ctx.task)
+
+        handler = call_core
+        for interceptor in reversed(self.interceptors):
+            next_handler = handler
+
+            async def call_interceptor(
+                ctx: WorkflowTaskInterceptorContext,
+                *,
+                interceptor: WorkerInterceptor = interceptor,
+                next_handler: Callable[
+                    [WorkflowTaskInterceptorContext],
+                    Awaitable[list[dict[str, Any]] | None],
+                ] = next_handler,
+            ) -> list[dict[str, Any]] | None:
+                return await interceptor.execute_workflow_task(ctx, next_handler)
+
+            handler = call_interceptor
+
+        return await handler(context)
+
+    async def _run_workflow_task_core(self, task: dict[str, Any]) -> list[dict[str, Any]] | None:
         import json as _json
 
         log.debug("workflow task payload: %s", _json.dumps(task, default=str)[:2000])
@@ -446,7 +483,7 @@ async def _run_activity_task(self, task: dict[str, Any]) -> str:
         )
         _set_context(act_ctx)
         try:
-            result = fn(*args) if not asyncio.iscoroutinefunction(fn) else await fn(*args)
+            result = await self._execute_activity_callable(task, activity_type, tuple(args), fn)
         except ActivityCancelled:
             log.info("activity %s cancelled via heartbeat", task_id)
             try:
@@ -508,7 +545,70 @@ async def _run_activity_task(self, task: dict[str, Any]) -> str:
             return "complete_error"
         return "completed"
 
+    async def _execute_activity_callable(
+        self,
+        task: dict[str, Any],
+        activity_type: str,
+        args: tuple[Any, ...],
+        fn: Callable[..., Any],
+    ) -> Any:
+        context = ActivityInterceptorContext(
+            worker_id=self.worker_id,
+            task_queue=self.task_queue,
+            task=task,
+            activity_type=activity_type,
+            args=args,
+        )
+
+        async def call_activity(ctx: ActivityInterceptorContext) -> Any:
+            result = fn(*ctx.args)
+            if asyncio.iscoroutine(result):
+                return await result
+            return result
+
+        handler = call_activity
+        for interceptor in reversed(self.interceptors):
+            next_handler = handler
+
+            async def call_interceptor(
+                ctx: ActivityInterceptorContext,
+                *,
+                interceptor: WorkerInterceptor = interceptor,
+                next_handler: Callable[[ActivityInterceptorContext], Awaitable[Any]] = next_handler,
+            ) -> Any:
+                return await interceptor.execute_activity(ctx, next_handler)
+
+            handler = call_interceptor
+
+        return await handler(context)
+
     async def _run_query_task(self, task: dict[str, Any]) -> str:
+        context = QueryTaskInterceptorContext(
+            worker_id=self.worker_id,
+            task_queue=self.task_queue,
+            task=task,
+        )
+
+        async def call_core(ctx: QueryTaskInterceptorContext) -> str:
+            return await self._run_query_task_core(ctx.task)
+
+        handler = call_core
+        for interceptor in reversed(self.interceptors):
+            next_handler = handler
+
+            async def call_interceptor(
+                ctx: QueryTaskInterceptorContext,
+                *,
+                interceptor: WorkerInterceptor = interceptor,
+                next_handler: Callable[[QueryTaskInterceptorContext], Awaitable[str]] = next_handler,
+            ) -> str:
+                return await interceptor.execute_query_task(ctx, next_handler)
+
+            handler = call_interceptor
+
+        return await handler(context)
+
+    async def _run_query_task_core(self, task: dict[str, Any]) -> str:
         query_task_id: str = task["query_task_id"]
         attempt: int = task.get("query_task_attempt", 1)
         wf_type: str = task.get("workflow_type", "")