From f89cc0721f3b09a21afbfa1c4e424d139d492395 Mon Sep 17 00:00:00 2001 From: Rajan-vish Date: Fri, 6 Nov 2026 16:41:23 +0530 Subject: [PATCH 1/2] fix(middleware): prevent smart retry delay compounding without schedule_source (#716) --- taskiq/middlewares/smart_retry_middleware.py | 25 +- tests/middlewares/test_smart_retry.py | 238 +++++++++++++++++++ 2 files changed, 261 insertions(+), 2 deletions(-) create mode 100644 tests/middlewares/test_smart_retry.py diff --git a/taskiq/middlewares/smart_retry_middleware.py b/taskiq/middlewares/smart_retry_middleware.py index 9752b07f..40693d98 100644 --- a/taskiq/middlewares/smart_retry_middleware.py +++ b/taskiq/middlewares/smart_retry_middleware.py @@ -106,6 +106,24 @@ def is_retry_on_error(self, message: TaskiqMessage) -> bool: retry_on_error = self.default_retry_label return retry_on_error + def _get_initial_delay(self, message: TaskiqMessage) -> float: + """ + Resolve the initial base delay for a task. + + If a private `_initial_delay` label was recorded on a previous retry, + it is reused to prevent compounding delays. Otherwise, it falls back + to the task's `delay` label or the middleware's `default_delay`. + + :param message: Task message. + :return: Initial delay in seconds. + """ + initial_delay = message.labels.get("_initial_delay") + if initial_delay is None: + initial_delay = message.labels.get("delay") + if initial_delay is None: + initial_delay = self.default_delay + return float(initial_delay) + def make_delay(self, message: TaskiqMessage, retries: int) -> float: """ Calculate retry delay. @@ -116,7 +134,7 @@ def make_delay(self, message: TaskiqMessage, retries: int) -> float: :param retries: Current retry count. :return: Delay in seconds. """ - delay = float(message.labels.get("delay", self.default_delay)) + delay = self._get_initial_delay(message) if self.use_delay_exponent: delay = min(delay * retries, self.max_delay_exponent) @@ -200,7 +218,10 @@ async def on_error( labels=message.labels, ) .with_task_id(message.task_id) - .with_labels(_retries=retries) + .with_labels( + _retries=retries, + _initial_delay=self._get_initial_delay(message), + ) ) await self.on_send(kicker, message, delay) diff --git a/tests/middlewares/test_smart_retry.py b/tests/middlewares/test_smart_retry.py new file mode 100644 index 00000000..7d6e4b45 --- /dev/null +++ b/tests/middlewares/test_smart_retry.py @@ -0,0 +1,238 @@ +import datetime +import uuid +from unittest.mock import AsyncMock + +import pytest + +from taskiq import InMemoryBroker, ScheduleSource, ScheduledTask, SmartRetryMiddleware +from taskiq.formatters.json_formatter import JSONFormatter +from taskiq.message import TaskiqMessage +from taskiq.result import TaskiqResult + + +@pytest.fixture +def broker() -> AsyncMock: + """Create a mock broker for testing.""" + 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 + + +class MockScheduleSource(ScheduleSource): + """Mock schedule source recording last scheduled task.""" + + def __init__(self) -> None: + self.last: ScheduledTask | None = None + + async def get_schedules(self) -> list[ScheduledTask]: + return [] + + async def add_schedule(self, schedule: ScheduledTask) -> None: + self.last = schedule + + +async def test_smart_retry_delay_exponent_without_schedule_source() -> None: + """Ensure retry delay does not compound without schedule_source.""" + mw = SmartRetryMiddleware( + default_retry_label=True, + default_retry_count=6, + default_delay=5.0, + use_delay_exponent=True, + max_delay_exponent=60.0, + schedule_source=None, + ) + broker = InMemoryBroker().with_middlewares(mw) + last_message: TaskiqMessage | None = None + + async def mock_kick(message: TaskiqMessage) -> None: + nonlocal last_message + last_message = message + + broker.kick = mock_kick # type: ignore[assignment] + + labels: dict[str, str | float] = {} + delays: list[float] = [] + + for _ in range(5): + msg = TaskiqMessage( + task_id="test-task", + task_name="test-task", + labels=labels, + args=[], + kwargs={}, + ) + result: TaskiqResult[None] = TaskiqResult( + is_err=True, + return_value=None, + execution_time=0.0, + ) + await mw.on_error(msg, result, ValueError("retry me")) + assert last_message is not None + labels = dict(last_message.labels) + delays.append(round(float(labels["delay"]), 1)) + + # Delay should scale linearly (5, 10, 15, 20, 25), not compound (5, 10, 30, ...) + assert delays == [5.0, 10.0, 15.0, 20.0, 25.0] + + +async def test_smart_retry_consistency_with_and_without_schedule_source() -> None: + """Verify delays match between broker path and schedule_source path.""" + source = MockScheduleSource() + mw_broker = SmartRetryMiddleware( + default_retry_label=True, + default_retry_count=6, + default_delay=5.0, + use_delay_exponent=True, + max_delay_exponent=60.0, + schedule_source=None, + ) + mw_source = SmartRetryMiddleware( + default_retry_label=True, + default_retry_count=6, + default_delay=5.0, + use_delay_exponent=True, + max_delay_exponent=60.0, + schedule_source=source, + ) + + broker_for_direct = InMemoryBroker().with_middlewares(mw_broker) + last_direct_message: TaskiqMessage | None = None + + async def mock_kick(message: TaskiqMessage) -> None: + nonlocal last_direct_message + last_direct_message = message + + broker_for_direct.kick = mock_kick # type: ignore[assignment] + InMemoryBroker().with_middlewares(mw_source) + + direct_labels: dict[str, str | float] = {} + source_labels: dict[str, str | float] = {} + direct_delays: list[float] = [] + source_delays: list[float] = [] + + for _ in range(4): + now = datetime.datetime.now(datetime.timezone.utc) + # Broker path + msg_direct = TaskiqMessage( + task_id="t1", + task_name="t1", + labels=direct_labels, + args=[], + kwargs={}, + ) + res_direct: TaskiqResult[None] = TaskiqResult( + is_err=True, + return_value=None, + execution_time=0.0, + ) + await mw_broker.on_error(msg_direct, res_direct, ValueError()) + assert last_direct_message is not None + direct_labels = dict(last_direct_message.labels) + direct_delays.append(round(float(direct_labels["delay"]), 1)) + + # Schedule source path + msg_source = TaskiqMessage( + task_id="t2", + task_name="t2", + labels=source_labels, + args=[], + kwargs={}, + ) + res_source: TaskiqResult[None] = TaskiqResult( + is_err=True, + return_value=None, + execution_time=0.0, + ) + await mw_source.on_error(msg_source, res_source, ValueError()) + assert source.last is not None and source.last.time is not None + source_labels = dict(source.last.labels) + source_delays.append( + round((source.last.time - now).total_seconds(), 1), + ) + + assert direct_delays == source_delays == [5.0, 10.0, 15.0, 20.0] + + +async def test_smart_retry_jitter_does_not_accumulate() -> None: + """Verify jitter is not compounded into the base delay on each retry.""" + mw = SmartRetryMiddleware( + default_retry_label=True, + default_retry_count=5, + default_delay=10.0, + use_jitter=True, + use_delay_exponent=False, + schedule_source=None, + ) + broker = InMemoryBroker().with_middlewares(mw) + last_message: TaskiqMessage | None = None + + async def mock_kick(message: TaskiqMessage) -> None: + nonlocal last_message + last_message = message + + broker.kick = mock_kick # type: ignore[assignment] + + labels: dict[str, str | float] = {} + for _ in range(4): + msg = TaskiqMessage( + task_id="t", + task_name="t", + labels=labels, + args=[], + kwargs={}, + ) + result: TaskiqResult[None] = TaskiqResult( + is_err=True, + return_value=None, + execution_time=0.0, + ) + await mw.on_error(msg, result, ValueError()) + assert last_message is not None + labels = dict(last_message.labels) + delay = float(labels["delay"]) + # Delay should be base 10.0 + random(0, 1), so always strictly < 11.0 + assert 10.0 <= delay <= 11.0 + + +async def test_smart_retry_custom_delay_label() -> None: + """Verify custom delay label on message is preserved as initial delay.""" + mw = SmartRetryMiddleware( + default_retry_label=True, + default_retry_count=5, + default_delay=5.0, + use_delay_exponent=True, + max_delay_exponent=100.0, + ) + broker = InMemoryBroker().with_middlewares(mw) + last_message: TaskiqMessage | None = None + + async def mock_kick(message: TaskiqMessage) -> None: + nonlocal last_message + last_message = message + + broker.kick = mock_kick # type: ignore[assignment] + + labels: dict[str, str | float] = {"delay": 12.0} + delays: list[float] = [] + + for _ in range(3): + msg = TaskiqMessage( + task_id="t", + task_name="t", + labels=labels, + args=[], + kwargs={}, + ) + result: TaskiqResult[None] = TaskiqResult( + is_err=True, + return_value=None, + execution_time=0.0, + ) + await mw.on_error(msg, result, ValueError()) + assert last_message is not None + labels = dict(last_message.labels) + delays.append(round(float(labels["delay"]), 1)) + + assert delays == [12.0, 24.0, 36.0] From 206049472511be68fa6aedf2313e71beeaa2412d Mon Sep 17 00:00:00 2001 From: Rajan-vish Date: Fri, 6 Nov 2026 22:12:57 +0530 Subject: [PATCH 2/2] style: sort imports in test_smart_retry.py per pre-commit --- tests/middlewares/test_smart_retry.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/middlewares/test_smart_retry.py b/tests/middlewares/test_smart_retry.py index 7d6e4b45..ee6b7e9d 100644 --- a/tests/middlewares/test_smart_retry.py +++ b/tests/middlewares/test_smart_retry.py @@ -4,7 +4,7 @@ import pytest -from taskiq import InMemoryBroker, ScheduleSource, ScheduledTask, SmartRetryMiddleware +from taskiq import InMemoryBroker, ScheduledTask, ScheduleSource, SmartRetryMiddleware from taskiq.formatters.json_formatter import JSONFormatter from taskiq.message import TaskiqMessage from taskiq.result import TaskiqResult