From ab3db6d671f1635680785b7e68689c4a12fb805b Mon Sep 17 00:00:00 2001 From: "sergei.romanchuk" Date: Sun, 12 Jul 2026 22:57:29 +0200 Subject: [PATCH] fix(scheduler): preserve custom task IDs for interval schedules Forward explicit task IDs to interval ScheduledTask payloads, matching cron and time scheduling semantics. Add regression coverage for recurring IDs, generated dispatch IDs, and schedule serialization. --- docs/guide/scheduling-tasks.md | 4 + taskiq/kicker.py | 1 + tests/scheduler/test_scheduling_identity.py | 124 ++++++++++++++++++++ 3 files changed, 129 insertions(+) create mode 100644 tests/scheduler/test_scheduling_identity.py diff --git a/docs/guide/scheduling-tasks.md b/docs/guide/scheduling-tasks.md index 44e646f2..a51acd72 100644 --- a/docs/guide/scheduling-tasks.md +++ b/docs/guide/scheduling-tasks.md @@ -159,6 +159,10 @@ If you want to pass additional labels, you can call these methods on the `Kicker ) ``` +An explicit `.with_task_id(...)` is stored for cron, interval and time +schedules and is reused for every dispatch of a recurring schedule. Without an +explicit task id, Taskiq generates a new invocation id for each dispatch. + ::: warning Cool warning! The `with_broker` method won't do anything in this case, since we have a broker assigned to each scheduler. diff --git a/taskiq/kicker.py b/taskiq/kicker.py index dc113a7e..b8b9c073 100644 --- a/taskiq/kicker.py +++ b/taskiq/kicker.py @@ -242,6 +242,7 @@ async def schedule_by_interval( labels=message.labels, args=message.args, kwargs=message.kwargs, + task_id=self.custom_task_id, interval=interval, ) await source.add_schedule(scheduled) diff --git a/tests/scheduler/test_scheduling_identity.py b/tests/scheduler/test_scheduling_identity.py new file mode 100644 index 00000000..2aabf4f3 --- /dev/null +++ b/tests/scheduler/test_scheduling_identity.py @@ -0,0 +1,124 @@ +from datetime import datetime, timezone +from typing import Any, Literal + +import pytest + +from taskiq.abc.schedule_source import ScheduleSource +from taskiq.compat import model_dump +from taskiq.kicker import AsyncKicker +from taskiq.message import BrokerMessage +from taskiq.scheduler.scheduled_task import ScheduledTask +from taskiq.scheduler.scheduler import TaskiqScheduler +from tests.utils import AsyncQueueBroker + +ScheduleKind = Literal["cron", "interval", "time"] + + +class RecordingBroker(AsyncQueueBroker): + """Record scheduler dispatch without using a transport.""" + + def __init__(self) -> None: + super().__init__() + self.sent: list[BrokerMessage] = [] + + async def kick(self, message: BrokerMessage) -> None: + self.sent.append(message) + + +class RecordingScheduleSource(ScheduleSource): + """Store schedules created through the public kicker API.""" + + def __init__(self) -> None: + self.schedules: list[ScheduledTask] = [] + + async def get_schedules(self) -> list[ScheduledTask]: + return self.schedules + + async def add_schedule(self, schedule: ScheduledTask) -> None: + self.schedules.append(schedule) + + +async def create_schedule( + kind: ScheduleKind, + kicker: AsyncKicker[Any, Any], + source: RecordingScheduleSource, +) -> ScheduledTask: + """Create one schedule through the selected public kicker method.""" + if kind == "cron": + created = await kicker.schedule_by_cron(source, "* * * * *") + elif kind == "interval": + created = await kicker.schedule_by_interval(source, 10) + else: + created = await kicker.schedule_by_time( + source, + datetime(2030, 1, 1, tzinfo=timezone.utc), + ) + return created.task + + +@pytest.mark.parametrize("kind", ["cron", "interval", "time"]) +async def test_all_schedule_methods_preserve_custom_task_id( + kind: ScheduleKind, +) -> None: + broker = RecordingBroker() + source = RecordingScheduleSource() + kicker: AsyncKicker[Any, Any] = AsyncKicker("demo.task", broker, {}) + + scheduled = await create_schedule( + kind, + kicker.with_task_id("custom-task-id"), + source, + ) + + assert scheduled.task_id == "custom-task-id" + assert source.schedules == [scheduled] + + +async def test_interval_without_custom_task_id_keeps_deferred_generation() -> None: + broker = RecordingBroker() + source = RecordingScheduleSource() + kicker: AsyncKicker[Any, Any] = AsyncKicker("demo.task", broker, {}) + + scheduled = await create_schedule("interval", kicker, source) + + assert scheduled.task_id is None + + +async def test_scheduler_reuses_stored_interval_task_id() -> None: + broker = RecordingBroker() + source = RecordingScheduleSource() + kicker: AsyncKicker[Any, Any] = AsyncKicker("demo.task", broker, {}) + scheduled = await create_schedule( + "interval", + kicker.with_task_id("custom-task-id"), + source, + ) + restored = ScheduledTask(**model_dump(scheduled)) + + scheduler = TaskiqScheduler(broker, [source]) + await scheduler.on_ready(source, restored) + await scheduler.on_ready(source, restored) + + assert [message.task_id for message in broker.sent] == [ + "custom-task-id", + "custom-task-id", + ] + + +async def test_scheduler_generates_task_id_for_each_interval_dispatch() -> None: + broker = RecordingBroker() + source = RecordingScheduleSource() + kicker: AsyncKicker[Any, Any] = AsyncKicker("demo.task", broker, {}) + scheduled = await create_schedule("interval", kicker, source) + restored = ScheduledTask(**model_dump(scheduled)) + generated_task_ids = iter(("generated-task-id-1", "generated-task-id-2")) + broker.with_id_generator(lambda: next(generated_task_ids)) + + scheduler = TaskiqScheduler(broker, [source]) + await scheduler.on_ready(source, restored) + await scheduler.on_ready(source, restored) + + assert [message.task_id for message in broker.sent] == [ + "generated-task-id-1", + "generated-task-id-2", + ]