Add acknowledgement progress support

docs/extending-taskiq/broker.md

@@ -40,9 +40,16 @@ async def listen(self) -> AsyncGenerator[AckableMessage, None]:
          # So you either set here method of a message,
          # or you can make a closure.
          ack=message.ack,
+         # Optionally expose a callback that reports the message
+         # is still being processed without finally acknowledging it.
+         ack_progress=message.ack_progress,
       )
 ```
 
+Tasks can call the optional progress callback through `Context.ack_progress()`.
+The callback may be invoked multiple times and should extend or reset the broker's
+acknowledgement deadline without marking the message as finally acknowledged.
+
 ## Conventions
 
 For brokers, we have several conventions. It's good if your broker implements them.

docs/guide/cli.md

@@ -74,12 +74,19 @@ async def critical_task() -> None:
 
 @broker.task(ack_type="manual")
 async def manually_acked_task(context: Context = TaskiqDepends()) -> None:
-    ...
+    while still_working:
+        await do_some_work()
+        if context.is_ack_progressable:
+            await context.ack_progress()
     await context.ack()
 ```
 
 If a task has no `ack_type` label, the worker-level `--ack-type` value is used.
 Manual acknowledgement requires a broker that yields `AckableMessage`.
+Some brokers can also report that a task is still being processed through
+`Context.ack_progress()` to prevent acknowledgement timeouts during long-running tasks.
+Currently, this is supported only by the `PushBasedJetStreamBroker` and
+`PullBasedJetStreamBroker` implementations from `taskiq-nats`.
 
 ### Type casts
 

taskiq/acks.py

@@ -55,20 +55,31 @@ class AckableMessage(BaseModel):
 
     data: bytes
     ack: Callable[[], None | Awaitable[None]]
+    ack_progress: Callable[[], None | Awaitable[None]] | None = None
 
 
 class AckController:
     """Controls acknowledgement state for a received message."""
 
-    def __init__(self, ack: Callable[[], None | Awaitable[None]] | None) -> None:
+    def __init__(
+        self,
+        ack: Callable[[], None | Awaitable[None]] | None,
+        ack_progress: Callable[[], None | Awaitable[None]] | None = None,
+    ) -> None:
         self._ack = ack
+        self._ack_progress = ack_progress
         self.is_acked = False
 
     @property
     def is_ackable(self) -> bool:
         """Whether the current message supports acknowledgement."""
         return self._ack is not None
 
+    @property
+    def is_ack_progressable(self) -> bool:
+        """Whether the current message supports acknowledgement progress."""
+        return self._ack_progress is not None
+
     async def ack(self) -> None:
         """Acknowledge the current message once."""
         if self._ack is None:
@@ -77,3 +88,11 @@ async def ack(self) -> None:
             return
         await maybe_awaitable(self._ack())
         self.is_acked = True
+
+    async def ack_progress(self) -> None:
+        """Report that the current message is still being processed."""
+        if self._ack_progress is None:
+            raise RuntimeError("Current message does not support ack progress.")
+        if self.is_acked:
+            return
+        await maybe_awaitable(self._ack_progress())

taskiq/context.py

@@ -34,6 +34,14 @@ def is_acked(self) -> bool:
         """Whether the current message has already been acknowledged."""
         return self._ack_controller is not None and self._ack_controller.is_acked
 
+    @property
+    def is_ack_progressable(self) -> bool:
+        """Whether the current message supports acknowledgement progress."""
+        return (
+            self._ack_controller is not None
+            and self._ack_controller.is_ack_progressable
+        )
+
     async def ack(self) -> None:
         """
         Acknowledge current message.
@@ -44,6 +52,16 @@ async def ack(self) -> None:
             raise RuntimeError("Current message is not ackable.")
         await self._ack_controller.ack()
 
+    async def ack_progress(self) -> None:
+        """
+        Report that the current message is still being processed.
+
+        :raises RuntimeError: if acknowledgement progress is not supported.
+        """
+        if self._ack_controller is None:
+            raise RuntimeError("Current message does not support ack progress.")
+        await self._ack_controller.ack_progress()
+
     async def requeue(self) -> None:
         """
         Requeue task.

taskiq/receiver/receiver.py

@@ -118,6 +118,7 @@ async def callback(  # noqa: C901, PLR0912
         """
         ack_controller = AckController(
             message.ack if isinstance(message, AckableMessage) else None,
+            message.ack_progress if isinstance(message, AckableMessage) else None,
         )
         message_data = message.data if isinstance(message, AckableMessage) else message
         try:

tests/receiver/test_receiver.py

@@ -539,6 +539,7 @@ async def test_manual_task_ack_from_context() -> None:
     async def my_task(context: Context = Depends()) -> int:
         events.append("task")
         assert context.is_ackable
+        assert not context.is_ack_progressable
         assert not context.is_acked
         await context.ack()
         assert context.is_acked
@@ -560,6 +561,55 @@ def ack_callback() -> None:
     assert events == ["task", "ack", "post_execute", "save", "post_save"]
 
 
+async def test_task_reports_ack_progress_from_context() -> None:
+    """A task can repeatedly report acknowledgement progress through context."""
+    events: list[str] = []
+    broker = (
+        InMemoryBroker()
+        .with_result_backend(
+            _EventResultBackend(events),
+        )
+        .with_middlewares(_EventMiddleware(events))
+    )
+
+    @broker.task(ack_type="manual")
+    async def my_task(context: Context = Depends()) -> int:
+        events.append("task")
+        assert context.is_ack_progressable
+        await context.ack_progress()
+        await context.ack_progress()
+        await context.ack()
+        await context.ack_progress()
+        return 1
+
+    def ack_callback() -> None:
+        events.append("ack")
+
+    def ack_progress_callback() -> None:
+        events.append("progress")
+
+    receiver = get_receiver(broker)
+    broker_message = broker.formatter.dumps(my_task.kicker()._prepare_message())
+
+    await receiver.callback(
+        AckableMessage(
+            data=broker_message.message,
+            ack=ack_callback,
+            ack_progress=ack_progress_callback,
+        ),
+    )
+
+    assert events == [
+        "task",
+        "progress",
+        "progress",
+        "ack",
+        "post_execute",
+        "save",
+        "post_save",
+    ]
+
+
 async def test_manual_task_ack_is_idempotent() -> None:
     """Calling Context.ack twice acknowledges the message once."""
     events: list[str] = []

tests/test_acks.py

@@ -1,6 +1,37 @@
 import pytest
 
-from taskiq.acks import AcknowledgeType, parse_acknowledge_type
+from taskiq.acks import AckController, AcknowledgeType, parse_acknowledge_type
+
+
+async def test_ack_progress_can_be_reported_repeatedly() -> None:
+    events: list[str] = []
+
+    def ack_callback() -> None:
+        events.append("ack")
+
+    async def ack_progress_callback() -> None:
+        events.append("progress")
+
+    controller = AckController(ack_callback, ack_progress_callback)
+
+    assert controller.is_ackable
+    assert controller.is_ack_progressable
+    await controller.ack_progress()
+    await controller.ack_progress()
+    assert not controller.is_acked
+    await controller.ack()
+    await controller.ack_progress()
+
+    assert controller.is_acked
+    assert events == ["progress", "progress", "ack"]
+
+
+async def test_ack_progress_raises_when_unsupported() -> None:
+    controller = AckController(lambda: None)
+
+    assert not controller.is_ack_progressable
+    with pytest.raises(RuntimeError, match="does not support ack progress"):
+        await controller.ack_progress()
 
 
 def test_parse_acknowledge_type_from_enum() -> None: