Add FailureHandler and CompletionHandler base classes for dependency hooks (#299)

## Summary

Introduces two new abstract base classes that let dependencies control
post-execution flow:

- `FailureHandler` - controls what happens when a task fails. `Retry`
now inherits from this and implements `handle_failure()` to schedule
retries.

- `CompletionHandler` - controls what happens after task completion.
`Perpetual` now inherits from this and implements `on_complete()` to
schedule the next execution.

This follows the pattern established by `Runtime` for `Timeout`: the
Worker delegates control to specialized dependencies rather than knowing
the details itself. The Worker's `_retry_if_requested()` and
`_perpetuate_if_requested()` methods are removed (~50 lines), with that
logic now living in the dependencies.

The dependency hierarchy now looks like:

```
Dependency (base - can observe via __aexit__)
├── Runtime (controls HOW task executes - Timeout)
├── FailureHandler (controls WHAT HAPPENS on failure - Retry)
└── CompletionHandler (controls WHAT HAPPENS after completion - Perpetual)
```

All three have `single = True` because only one thing can control each
aspect.

Also adds `after(timedelta)` and `at(datetime)` methods to both
`Perpetual` and `Retry`, giving developers flexibility to schedule the
next execution using whatever they have handy. For `Retry`, the old
`in_()` method is kept as a backwards-compatible alias.

Closes #297

🤖 Generated with [Claude Code](https://claude.ai/code)

---------

Co-authored-by: Claude Opus 4.5 <noreply@anthropic.com>

Author: chrisguidry

loq.toml

@@ -10,20 +10,20 @@ max_lines = 750
 # Source files that still need exceptions above 750
 [[rules]]
 path = "src/docket/worker.py"
-max_lines = 1150
+max_lines = 1125
 
 [[rules]]
 path = "src/docket/cli.py"
 max_lines = 945
 
 [[rules]]
 path = "src/docket/execution.py"
-max_lines = 900
+max_lines = 903
 
 [[rules]]
 path = "src/docket/docket.py"
-max_lines = 900
+max_lines = 866
 
 [[rules]]
 path = "src/docket/strikelist.py"
-max_lines = 650
+max_lines = 616

src/docket/__init__.py

@@ -27,7 +27,7 @@
     Timeout,
 )
 from .docket import Docket
-from .execution import Execution, ExecutionState
+from .execution import Execution, ExecutionCancelled, ExecutionState
 from .strikelist import StrikeList
 from .worker import Worker
 from . import testing
@@ -42,6 +42,7 @@
     "Depends",
     "Docket",
     "Execution",
+    "ExecutionCancelled",
     "ExecutionState",
     "ExponentialRetry",
     "Logged",

src/docket/dependencies/__init__.py

@@ -6,7 +6,15 @@
 
 from __future__ import annotations
 
-from ._base import AdmissionBlocked, Dependency, Runtime
+from ._base import (
+    AdmissionBlocked,
+    CompletionHandler,
+    Dependency,
+    FailureHandler,
+    Runtime,
+    TaskOutcome,
+    format_duration,
+)
 from ._concurrency import ConcurrencyBlocked, ConcurrencyLimit
 from ._contextual import (
     CurrentDocket,
@@ -41,6 +49,10 @@
     # Base
     "Dependency",
     "Runtime",
+    "FailureHandler",
+    "CompletionHandler",
+    "TaskOutcome",
+    "format_duration",
     # Contextual dependencies
     "CurrentDocket",
     "CurrentExecution",

src/docket/dependencies/_base.py

@@ -4,6 +4,8 @@
 
 import abc
 from contextvars import ContextVar
+from dataclasses import dataclass, field
+from datetime import timedelta
 from types import TracebackType
 from typing import TYPE_CHECKING, Any, Awaitable, Callable
 
@@ -13,6 +15,23 @@
     from ..worker import Worker
 
 
+def format_duration(seconds: float) -> str:
+    """Format a duration for log output."""
+    if seconds < 100:
+        return f"{seconds * 1000:6.0f}ms"
+    else:
+        return f"{seconds:6.0f}s "
+
+
+@dataclass
+class TaskOutcome:
+    """Captures the outcome of a task execution for handlers."""
+
+    duration: timedelta
+    result: Any = field(default=None)
+    exception: BaseException | None = field(default=None)
+
+
 class AdmissionBlocked(Exception):
     """Raised when a task cannot start due to admission control.
 
@@ -72,3 +91,57 @@ async def run(
             kwargs: Keyword arguments including resolved dependencies
         """
         ...  # pragma: no cover
+
+
+class FailureHandler(Dependency):
+    """Base class for dependencies that control what happens when a task fails.
+
+    Called on exceptions. If handle_failure() returns True, the handler
+    took responsibility (e.g., scheduled a retry) and Worker won't mark
+    the execution as failed.
+
+    Only one FailureHandler per task (single=True).
+    """
+
+    single = True
+
+    @abc.abstractmethod
+    async def handle_failure(self, execution: Execution, outcome: TaskOutcome) -> bool:
+        """Handle a task failure.
+
+        Args:
+            execution: The task execution context
+            outcome: The task outcome containing duration and exception
+
+        Returns:
+            True if handled (Worker won't mark as failed)
+            False if not handled (Worker proceeds normally)
+        """
+        ...  # pragma: no cover
+
+
+class CompletionHandler(Dependency):
+    """Base class for dependencies that control what happens after task completion.
+
+    Called after execution is truly done (success, or failure with no retry).
+    If on_complete() returns True, the handler took responsibility (e.g.,
+    scheduled follow-up work) and did its own logging.
+
+    Only one CompletionHandler per task (single=True).
+    """
+
+    single = True
+
+    @abc.abstractmethod
+    async def on_complete(self, execution: Execution, outcome: TaskOutcome) -> bool:
+        """Handle task completion.
+
+        Args:
+            execution: The task execution context
+            outcome: The task outcome containing duration, result, and exception
+
+        Returns:
+            True if handled (did own logging/metrics)
+            False if not handled (Worker does normal logging)
+        """
+        ...  # pragma: no cover

src/docket/dependencies/_perpetual.py

@@ -2,13 +2,21 @@
 
 from __future__ import annotations
 
-from datetime import timedelta
-from typing import Any
+import logging
+from datetime import datetime, timedelta, timezone
+from typing import TYPE_CHECKING, Any
 
-from ._base import Dependency
+from ._base import CompletionHandler, TaskOutcome, format_duration
 
+if TYPE_CHECKING:  # pragma: no cover
+    from ..execution import Execution
 
-class Perpetual(Dependency):
+from ..instrumentation import TASKS_PERPETUATED
+
+logger = logging.getLogger(__name__)
+
+
+class Perpetual(CompletionHandler):
     """Declare a task that should be run perpetually.  Perpetual tasks are automatically
     rescheduled for the future after they finish (whether they succeed or fail).  A
     perpetual task can be scheduled at worker startup with the `automatic=True`.
@@ -31,6 +39,7 @@ async def my_task(perpetual: Perpetual = Perpetual()) -> None:
     kwargs: dict[str, Any]
 
     cancelled: bool
+    _next_when: datetime | None
 
     def __init__(
         self,
@@ -48,6 +57,7 @@ def __init__(
         self.every = every
         self.automatic = automatic
         self.cancelled = False
+        self._next_when = None
 
     async def __aenter__(self) -> Perpetual:
         execution = self.execution.get()
@@ -62,3 +72,42 @@ def cancel(self) -> None:
     def perpetuate(self, *args: Any, **kwargs: Any) -> None:
         self.args = args
         self.kwargs = kwargs
+
+    def after(self, delay: timedelta) -> None:
+        """Schedule the next execution after the given delay."""
+        self._next_when = datetime.now(timezone.utc) + delay
+
+    def at(self, when: datetime) -> None:
+        """Schedule the next execution at the given time."""
+        self._next_when = when
+
+    async def on_complete(self, execution: Execution, outcome: TaskOutcome) -> bool:
+        """Handle completion by scheduling the next execution."""
+        if self.cancelled:
+            docket = self.docket.get()
+            async with docket.redis() as redis:
+                await docket._cancel(redis, execution.key)
+            return False
+
+        docket = self.docket.get()
+        worker = self.worker.get()
+
+        if self._next_when:
+            when = self._next_when
+        else:
+            now = datetime.now(timezone.utc)
+            when = max(now, now + self.every - outcome.duration)
+
+        await docket.replace(execution.function, when, execution.key)(
+            *self.args,
+            **self.kwargs,
+        )
+
+        TASKS_PERPETUATED.add(1, {**worker.labels(), **execution.general_labels()})
+        logger.info(
+            "↫ [%s] %s",
+            format_duration(outcome.duration.total_seconds()),
+            execution.call_repr(),
+        )
+
+        return True

src/docket/dependencies/_retry.py

@@ -2,17 +2,25 @@
 
 from __future__ import annotations
 
+import logging
 from datetime import datetime, timedelta, timezone
-from typing import NoReturn
+from typing import TYPE_CHECKING, NoReturn
 
-from ._base import Dependency
+from ._base import FailureHandler, TaskOutcome, format_duration
+
+if TYPE_CHECKING:  # pragma: no cover
+    from ..execution import Execution
+
+from ..instrumentation import TASKS_RETRIED
+
+logger = logging.getLogger(__name__)
 
 
 class ForcedRetry(Exception):
-    """Raised when a task requests a retry via `in_` or `at`"""
+    """Raised when a task requests a retry via `after` or `at`"""
 
 
-class Retry(Dependency):
+class Retry(FailureHandler):
     """Configures linear retries for a task.  You can specify the total number of
     attempts (or `None` to retry indefinitely), and the delay between attempts.
 
@@ -50,16 +58,40 @@ async def __aenter__(self) -> Retry:
         retry.attempt = execution.attempt
         return retry
 
+    def after(self, delay: timedelta) -> NoReturn:
+        """Request a retry after the given delay."""
+        self.delay = delay
+        raise ForcedRetry()
+
     def at(self, when: datetime) -> NoReturn:
+        """Request a retry at the given time."""
         now = datetime.now(timezone.utc)
         diff = when - now
         diff = diff if diff.total_seconds() >= 0 else timedelta(0)
+        self.after(diff)
+
+    def in_(self, delay: timedelta) -> NoReturn:
+        """Deprecated: use after() instead."""
+        self.after(delay)
+
+    async def handle_failure(self, execution: Execution, outcome: TaskOutcome) -> bool:
+        """Handle failure by scheduling a retry if attempts remain."""
+        if self.attempts is not None and execution.attempt >= self.attempts:
+            return False
+
+        execution.when = datetime.now(timezone.utc) + self.delay
+        execution.attempt += 1
+        await execution.schedule(replace=True)
+
+        worker = self.worker.get()
+        TASKS_RETRIED.add(1, {**worker.labels(), **execution.general_labels()})
+        logger.info(
+            "↫ [%s] %s",
+            format_duration(outcome.duration.total_seconds()),
+            execution.call_repr(),
+        )
 
-        self.in_(diff)
-
-    def in_(self, when: timedelta) -> NoReturn:
-        self.delay = when
-        raise ForcedRetry()
+        return True
 
 
 class ExponentialRetry(Retry):

src/docket/execution.py

@@ -33,6 +33,13 @@
 
 logger: logging.Logger = logging.getLogger(__name__)
 
+
+class ExecutionCancelled(Exception):
+    """Raised when get_result() is called on a cancelled execution."""
+
+    pass
+
+
 TaskFunction = Callable[..., Awaitable[Any]]
 Message = dict[bytes, bytes]
 
@@ -682,8 +689,14 @@ async def get_result(
         if timeout is not None:
             deadline = datetime.now(timezone.utc) + timeout
 
+        terminal_states = (
+            ExecutionState.COMPLETED,
+            ExecutionState.FAILED,
+            ExecutionState.CANCELLED,
+        )
+
         # Wait for execution to complete if not already done
-        if self.state not in (ExecutionState.COMPLETED, ExecutionState.FAILED):
+        if self.state not in terminal_states:
             # Calculate timeout duration if absolute deadline provided
             timeout_seconds = None
             if deadline is not None:
@@ -701,10 +714,7 @@ async def wait_for_completion():
                     async for event in self.subscribe():  # pragma: no branch
                         if event["type"] == "state":
                             state = ExecutionState(event["state"])
-                            if state in (
-                                ExecutionState.COMPLETED,
-                                ExecutionState.FAILED,
-                            ):
+                            if state in terminal_states:
                                 # Sync to get latest data including result key
                                 await self.sync()
                                 break
@@ -716,6 +726,10 @@ async def wait_for_completion():
                     f"Timeout waiting for execution {self.key} to complete"
                 )
 
+        # If cancelled, raise ExecutionCancelled
+        if self.state == ExecutionState.CANCELLED:
+            raise ExecutionCancelled(f"Execution {self.key} was cancelled")
+
         # If failed, retrieve and raise the exception
         if self.state == ExecutionState.FAILED:
             if self.result_key: