feat: add per-task types_of_exceptions (#651)

Author: spikeninja

docs/available-components/middlewares.md

@@ -34,6 +34,25 @@ async def test():
 
 `retry_on_error` enables retries for a task. `max_retries` is the maximum number of retry attempts.
 
+### Retrying only specific exceptions
+
+By default, all exceptions trigger a retry. You can limit retries to specific
+exception types either broker-wide via `types_of_exceptions`, or per task via
+the `types_of_exceptions` label in the task decorator. The per-task value
+overrides the broker-wide setting.
+
+```python
+broker = ZeroMQBroker().with_middlewares(
+    # Broker-wide default: retry only on ConnectionError.
+    SimpleRetryMiddleware(types_of_exceptions=(ConnectionError,)),
+)
+
+
+@broker.task(retry_on_error=True, types_of_exceptions=(ValueError, KeyError))
+async def test():
+    raise ValueError("retry only on ValueError or KeyError")
+```
+
 ## Smart retry middleware
 
 The `SmartRetryMiddleware` automatically retries tasks with flexible delay settings and retry strategies when errors occur. This is particularly useful when tasks fail due to temporary issues, such as network errors or temporary unavailability of external services.
@@ -78,6 +97,13 @@ async def my_task():
 * `retry_on_error`: Enables the retry mechanism for the specific task.
 * `max_retries`: Maximum number of retries (overrides middleware default).
 * `delay`: Initial delay before retrying the task, in seconds.
+* `types_of_exceptions`: Exception types that trigger a retry for this task. Overrides the broker-wide `types_of_exceptions` passed to the middleware.
+
+```python
+@broker.task(retry_on_error=True, types_of_exceptions=(ConnectionError,))
+async def my_task():
+    raise ConnectionError("retrying only on ConnectionError")
+```
 
 ### Usage Recommendations
 

taskiq/kicker.py

@@ -322,6 +322,12 @@ def _prepare_message(
             formatted_kwargs[kwarg_name] = self._prepare_arg(kwarg_val)
 
         for label, label_val in self.labels.items():
+            # `types_of_exceptions` is only ever read back from the
+            # locally registered task's labels (see retry middlewares),
+            # never from the wire. It holds exception classes, which
+            # can't be faithfully serialized, so don't ship it at all.
+            if label == "types_of_exceptions":
+                continue
             labels[label], labels_types[label] = prepare_label(label_val)
 
         task_id = self.custom_task_id

taskiq/middlewares/simple_retry_middleware.py

@@ -26,6 +26,28 @@ def __init__(
         self.no_result_on_retry = no_result_on_retry
         self.types_of_exceptions = types_of_exceptions
 
+    def _get_types_of_exceptions(
+        self,
+        message: "TaskiqMessage",
+    ) -> Iterable[type[BaseException]] | None:
+        """
+        Resolve retryable exception types for a task.
+
+        Per-task ``types_of_exceptions`` set via the task decorator take
+        precedence over the broker-wide value. Types are read from the
+        registered task object, since label values are stringified when a
+        message is serialized and cannot carry real exception types.
+
+        :param message: Original task message.
+        :return: Effective exception types or None.
+        """
+        task = self.broker.find_task(message.task_name)
+        if task is not None:
+            task_types = task.labels.get("types_of_exceptions")
+            if task_types is not None:
+                return task_types
+        return self.types_of_exceptions
+
     async def on_error(
         self,
         message: "TaskiqMessage",
@@ -45,9 +67,10 @@ async def on_error(
         :param result: execution result.
         :param exception: found exception.
         """
-        if self.types_of_exceptions is not None and not isinstance(
+        types_of_exceptions = self._get_types_of_exceptions(message)
+        if types_of_exceptions is not None and not isinstance(
             exception,
-            tuple(self.types_of_exceptions),
+            tuple(types_of_exceptions),
         ):
             return
 

taskiq/middlewares/smart_retry_middleware.py

@@ -68,6 +68,28 @@ def __init__(
                 "schedule_source must be an instance of ScheduleSource or None",
             )
 
+    def _get_types_of_exceptions(
+        self,
+        message: TaskiqMessage,
+    ) -> Iterable[type[BaseException]] | None:
+        """
+        Resolve retryable exception types for a task.
+
+        Per-task ``types_of_exceptions`` set via the task decorator take
+        precedence over the broker-wide value. Types are read from the
+        registered task object, since label values are stringified when a
+        message is serialized and cannot carry real exception types.
+
+        :param message: Original task message.
+        :return: Effective exception types or None.
+        """
+        task = self.broker.find_task(message.task_name)
+        if task is not None:
+            task_types = task.labels.get("types_of_exceptions")
+            if task_types is not None:
+                return task_types
+        return self.types_of_exceptions
+
     def is_retry_on_error(self, message: TaskiqMessage) -> bool:
         """
         Check if retry is enabled for this task.
@@ -142,9 +164,10 @@ async def on_error(
         :param result: Execution result.
         :param exception: Caught exception.
         """
-        if self.types_of_exceptions is not None and not isinstance(
+        types_of_exceptions = self._get_types_of_exceptions(message)
+        if types_of_exceptions is not None and not isinstance(
             exception,
-            tuple(self.types_of_exceptions),
+            tuple(types_of_exceptions),
         ):
             return
 

tests/middlewares/test_simple_retry.py

@@ -14,6 +14,7 @@ def broker() -> AsyncMock:
     mocked_broker = AsyncMock()
     mocked_broker.id_generator = lambda: uuid.uuid4().hex
     mocked_broker.formatter = JSONFormatter()
+    mocked_broker.find_task = lambda task_name: None
     return mocked_broker
 
 

tests/middlewares/test_task_retry.py

@@ -200,6 +200,122 @@ def run_task2() -> None:
     assert runs == 1
 
 
+@pytest.mark.parametrize(
+    "middleware_class",
+    [SimpleRetryMiddleware, SmartRetryMiddleware],
+)
+async def test_per_task_exc_types_not_matching(middleware_class: type) -> None:
+    # per-task types_of_exceptions does not include the raised exception
+    broker = InMemoryBroker().with_middlewares(
+        middleware_class(no_result_on_retry=True, default_retry_label=True),
+    )
+    runs = 0
+
+    @broker.task(max_retries=10, types_of_exceptions=(KeyError,))
+    def run_task() -> None:
+        nonlocal runs
+
+        runs += 1
+
+        raise ValueError(runs)
+
+    task = await run_task.kiq()
+    resp = await task.wait_result(timeout=1)
+    with pytest.raises(ValueError):
+        resp.raise_for_error()
+
+    assert runs == 1
+
+
+@pytest.mark.parametrize(
+    "middleware_class",
+    [SimpleRetryMiddleware, SmartRetryMiddleware],
+)
+async def test_per_task_exc_types_matching(middleware_class: type) -> None:
+    # per-task types_of_exceptions includes the raised exception
+    broker = InMemoryBroker().with_middlewares(
+        middleware_class(no_result_on_retry=True, default_retry_label=True),
+    )
+    runs = 0
+
+    @broker.task(max_retries=10, types_of_exceptions=(ValueError,))
+    def run_task() -> None:
+        nonlocal runs
+
+        runs += 1
+
+        raise ValueError(runs)
+
+    task = await run_task.kiq()
+    resp = await task.wait_result(timeout=1)
+    with pytest.raises(ValueError):
+        resp.raise_for_error()
+
+    assert runs == 10
+
+
+@pytest.mark.parametrize(
+    "middleware_class",
+    [SimpleRetryMiddleware, SmartRetryMiddleware],
+)
+async def test_per_task_exc_types_override_global(middleware_class: type) -> None:
+    # per-task types_of_exceptions takes precedence over broker-wide value
+    broker = InMemoryBroker().with_middlewares(
+        middleware_class(
+            no_result_on_retry=True,
+            default_retry_label=True,
+            types_of_exceptions=(KeyError,),
+        ),
+    )
+    runs = 0
+
+    @broker.task(max_retries=10, types_of_exceptions=(ValueError,))
+    def run_task() -> None:
+        nonlocal runs
+
+        runs += 1
+
+        raise ValueError(runs)
+
+    task = await run_task.kiq()
+    resp = await task.wait_result(timeout=1)
+    with pytest.raises(ValueError):
+        resp.raise_for_error()
+
+    assert runs == 10
+
+
+@pytest.mark.parametrize(
+    "middleware_class",
+    [SimpleRetryMiddleware, SmartRetryMiddleware],
+)
+async def test_global_exc_types_without_per_task(middleware_class: type) -> None:
+    # broker-wide types_of_exceptions still applies when no per-task value set
+    broker = InMemoryBroker().with_middlewares(
+        middleware_class(
+            no_result_on_retry=True,
+            default_retry_label=True,
+            types_of_exceptions=(KeyError,),
+        ),
+    )
+    runs = 0
+
+    @broker.task(max_retries=10)
+    def run_task() -> None:
+        nonlocal runs
+
+        runs += 1
+
+        raise ValueError(runs)
+
+    task = await run_task.kiq()
+    resp = await task.wait_result(timeout=1)
+    with pytest.raises(ValueError):
+        resp.raise_for_error()
+
+    assert runs == 1
+
+
 async def test_retry_of_custom_exc_types_of_smart_middleware() -> None:
     # test that the passed error will be handled
     broker = InMemoryBroker().with_middlewares(

tests/test_kicker.py

@@ -0,0 +1,45 @@
+from typing import Any
+
+from taskiq import InMemoryBroker
+from taskiq.kicker import AsyncKicker
+
+
+async def test_types_of_exceptions_not_serialized() -> None:
+    """`types_of_exceptions` label should never be sent over the wire."""
+    broker = InMemoryBroker()
+
+    @broker.task(types_of_exceptions=(ValueError, TypeError))
+    async def run_task() -> None:
+        pass
+
+    kicker = run_task.kicker()
+    message = kicker._prepare_message()
+
+    assert "types_of_exceptions" not in message.labels
+    assert "types_of_exceptions" not in (message.labels_types or {})
+
+
+async def test_types_of_exceptions_still_local_on_task() -> None:
+    """The registered task still keeps the real exception types locally."""
+    broker = InMemoryBroker()
+
+    @broker.task(types_of_exceptions=(ValueError, TypeError))
+    async def run_task() -> None:
+        pass
+
+    task = broker.find_task(run_task.task_name)
+    assert task is not None
+    assert task.labels["types_of_exceptions"] == (ValueError, TypeError)
+
+
+async def test_other_labels_still_serialized() -> None:
+    """Unrelated labels are unaffected by the fix."""
+    kicker: AsyncKicker[Any, Any] = AsyncKicker(
+        task_name="some_task",
+        broker=InMemoryBroker(),
+        labels={"retries": 3, "queue": "high_priority"},
+    )
+    message = kicker._prepare_message()
+
+    assert message.labels["retries"] == "3"
+    assert message.labels["queue"] == "high_priority"