Add manual task acknowledgement

Author: vvanglro

docs/guide/cli.md

@@ -47,10 +47,11 @@ We have two options for this:
 
 ### Acknowledgements
 
-The taskiq supports three types of acknowledgements:
+The taskiq supports four types of acknowledgements:
 * `when_received` - task is acknowledged when it is **received** by the worker.
 * `when_executed` - task is acknowledged right after it is **executed** by the worker.
 * `when_saved` - task is acknowledged when the result of execution is saved in the result backend.
+* `manual` - task is acknowledged by calling `Context.ack()` inside the task.
 
 This can be configured using `--ack-type` parameter. For example:
 
@@ -69,9 +70,16 @@ async def best_effort_task() -> None:
 @broker.task(ack_type="when_saved")
 async def critical_task() -> None:
     ...
+
+
+@broker.task(ack_type="manual")
+async def manually_acked_task(context: Context = TaskiqDepends()) -> None:
+    ...
+    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`.
 
 ### Type casts
 
@@ -159,7 +167,7 @@ The number of signals before a hard kill can be configured with the `--hardkill-
 * `--no-propagate-errors` - if this parameter is enabled, exceptions won't be thrown in generator dependencies.
 * `--receiver` - python path to custom receiver class.
 * `--receiver_arg` - custom args for receiver.
-* `--ack-type` - Type of acknowledgement. This parameter is used to set when to acknowledge the task. Possible values are `when_received`, `when_executed`, `when_saved`. Default is `when_saved`.
+* `--ack-type` - Type of acknowledgement. This parameter is used to set when to acknowledge the task. Possible values are `when_received`, `when_executed`, `when_saved`, `manual`. Default is `when_saved`.
 * `--max-tasks-per-child` - maximum number of tasks to be executed by a single worker process before restart.
 * `--max-fails` - Maximum number of child process exits.
 * `--shutdown-timeout` - maximum amount of time for graceful broker's shutdown in seconds (default 5).

taskiq/acks.py

@@ -4,6 +4,8 @@
 
 from pydantic import BaseModel
 
+from taskiq.utils import maybe_awaitable
+
 
 @enum.unique
 class AcknowledgeType(str, enum.Enum):
@@ -19,6 +21,9 @@ class AcknowledgeType(str, enum.Enum):
     # acknowledged when the task will be saved
     # only after it's saved in the result backend.
     WHEN_SAVED = "when_saved"
+    # This option means that the task is responsible
+    # for acknowledging the message through Context.ack.
+    MANUAL = "manual"
 
 
 def parse_acknowledge_type(value: Any) -> AcknowledgeType:
@@ -50,3 +55,25 @@ class AckableMessage(BaseModel):
 
     data: bytes
     ack: Callable[[], None | Awaitable[None]]
+
+
+class AckController:
+    """Controls acknowledgement state for a received message."""
+
+    def __init__(self, ack: Callable[[], None | Awaitable[None]] | None) -> None:
+        self._ack = ack
+        self.is_acked = False
+
+    @property
+    def is_ackable(self) -> bool:
+        """Whether the current message supports acknowledgement."""
+        return self._ack is not None
+
+    async def ack(self) -> None:
+        """Acknowledge the current message once."""
+        if self._ack is None:
+            raise RuntimeError("Current message is not ackable.")
+        if self.is_acked:
+            return
+        await maybe_awaitable(self._ack())
+        self.is_acked = True

taskiq/context.py

@@ -1,6 +1,7 @@
 from typing import TYPE_CHECKING
 
 from taskiq.abc.broker import AsyncBroker
+from taskiq.acks import AckController
 from taskiq.exceptions import NoResultError, TaskRejectedError
 from taskiq.message import TaskiqMessage
 
@@ -11,12 +12,38 @@
 class Context:
     """Context class."""
 
-    def __init__(self, message: TaskiqMessage, broker: AsyncBroker) -> None:
+    def __init__(
+        self,
+        message: TaskiqMessage,
+        broker: AsyncBroker,
+        ack_controller: AckController | None = None,
+    ) -> None:
         self.message = message
         self.broker = broker
+        self._ack_controller = ack_controller
         self.state: TaskiqState = None  # type: ignore
         self.state = broker.state
 
+    @property
+    def is_ackable(self) -> bool:
+        """Whether the current message supports acknowledgement."""
+        return self._ack_controller is not None and self._ack_controller.is_ackable
+
+    @property
+    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
+
+    async def ack(self) -> None:
+        """
+        Acknowledge current message.
+
+        :raises RuntimeError: if current broker message is not ackable.
+        """
+        if self._ack_controller is None:
+            raise RuntimeError("Current message is not ackable.")
+        await self._ack_controller.ack()
+
     async def requeue(self) -> None:
         """
         Requeue task.

taskiq/receiver/receiver.py

@@ -15,7 +15,7 @@
 
 from taskiq.abc.broker import AckableMessage, AsyncBroker
 from taskiq.abc.middleware import TaskiqMiddleware
-from taskiq.acks import AcknowledgeType, parse_acknowledge_type
+from taskiq.acks import AckController, AcknowledgeType, parse_acknowledge_type
 from taskiq.context import Context
 from taskiq.exceptions import NoResultError
 from taskiq.message import TaskiqMessage
@@ -116,6 +116,9 @@ async def callback(  # noqa: C901, PLR0912
         :param raise_err: raise an error if cannot save result in
             result_backend.
         """
+        ack_controller = AckController(
+            message.ack if isinstance(message, AckableMessage) else None,
+        )
         message_data = message.data if isinstance(message, AckableMessage) else message
         try:
             taskiq_msg = self.broker.formatter.loads(message=message_data)
@@ -155,22 +158,17 @@ async def callback(  # noqa: C901, PLR0912
             taskiq_msg.task_id,
         )
 
-        if ack_time == AcknowledgeType.WHEN_RECEIVED and isinstance(
-            message,
-            AckableMessage,
-        ):
-            await maybe_awaitable(message.ack())
+        if ack_time == AcknowledgeType.WHEN_RECEIVED and ack_controller.is_ackable:
+            await ack_controller.ack()
 
         result = await self.run_task(
             target=task.original_func,
             message=taskiq_msg,
+            ack_controller=ack_controller,
         )
 
-        if ack_time == AcknowledgeType.WHEN_EXECUTED and isinstance(
-            message,
-            AckableMessage,
-        ):
-            await maybe_awaitable(message.ack())
+        if ack_time == AcknowledgeType.WHEN_EXECUTED and ack_controller.is_ackable:
+            await ack_controller.ack()
 
         for middleware in reversed(self.broker.middlewares):
             if middleware.__class__.post_execute != TaskiqMiddleware.post_execute:
@@ -193,11 +191,8 @@ async def callback(  # noqa: C901, PLR0912
             if raise_err:
                 raise exc
 
-        if ack_time == AcknowledgeType.WHEN_SAVED and isinstance(
-            message,
-            AckableMessage,
-        ):
-            await maybe_awaitable(message.ack())
+        if ack_time == AcknowledgeType.WHEN_SAVED and ack_controller.is_ackable:
+            await ack_controller.ack()
 
     def _get_ack_time(self, message: TaskiqMessage) -> AcknowledgeType:
         """
@@ -219,6 +214,7 @@ async def run_task(  # noqa: C901, PLR0912, PLR0915
         self,
         target: Callable[..., Any],
         message: TaskiqMessage,
+        ack_controller: AckController | None = None,
     ) -> TaskiqResult[Any]:
         """
         This function actually executes functions.
@@ -259,7 +255,7 @@ async def run_task(  # noqa: C901, PLR0912, PLR0915
             broker_ctx = self.broker.custom_dependency_context
             broker_ctx.update(
                 {
-                    Context: Context(message, self.broker),
+                    Context: Context(message, self.broker, ack_controller),
                     TaskiqState: self.broker.state,
                 },
             )

tests/receiver/test_receiver.py

@@ -16,6 +16,7 @@
 from taskiq.abc.result_backend import AsyncResultBackend
 from taskiq.acks import AcknowledgeType
 from taskiq.brokers.inmemory_broker import InMemoryBroker
+from taskiq.context import Context
 from taskiq.exceptions import NoResultError, TaskiqResultTimeoutError
 from taskiq.message import TaskiqMessage
 from taskiq.receiver import Receiver
@@ -523,6 +524,112 @@ def ack_callback() -> None:
     assert events == ["ack", "task", "post_execute", "save", "post_save"]
 
 
+async def test_manual_task_ack_from_context() -> None:
+    """Manual ack lets a task acknowledge the message 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_ackable
+        assert not context.is_acked
+        await context.ack()
+        assert context.is_acked
+        return 1
+
+    def ack_callback() -> None:
+        events.append("ack")
+
+    receiver = get_receiver(broker, ack_type=AcknowledgeType.WHEN_SAVED)
+    broker_message = broker.formatter.dumps(my_task.kicker()._prepare_message())
+
+    await receiver.callback(
+        AckableMessage(
+            data=broker_message.message,
+            ack=ack_callback,
+        ),
+    )
+
+    assert events == ["task", "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] = []
+    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")
+        await context.ack()
+        await context.ack()
+        return 1
+
+    def ack_callback() -> None:
+        events.append("ack")
+
+    receiver = get_receiver(broker, ack_type=AcknowledgeType.WHEN_SAVED)
+    broker_message = broker.formatter.dumps(my_task.kicker()._prepare_message())
+
+    await receiver.callback(
+        AckableMessage(
+            data=broker_message.message,
+            ack=ack_callback,
+        ),
+    )
+
+    assert events == ["task", "ack", "post_execute", "save", "post_save"]
+
+
+async def test_manual_task_ack_without_ackable_message_saves_error() -> None:
+    """Context.ack fails clearly when broker messages do not support acking."""
+    events: list[str] = []
+    result_backend = _EventResultBackend(events)
+    broker = (
+        InMemoryBroker()
+        .with_result_backend(result_backend)
+        .with_middlewares(_EventMiddleware(events))
+    )
+
+    @broker.task(ack_type="manual")
+    async def my_task(context: Context = Depends()) -> int:
+        assert not context.is_ackable
+        await context.ack()
+        return 1
+
+    receiver = get_receiver(broker)
+    broker_message = broker.formatter.dumps(
+        TaskiqMessage(
+            task_id="task_id",
+            task_name=my_task.task_name,
+            labels={"ack_type": "manual"},
+            args=[],
+            kwargs={},
+        ),
+    )
+
+    await receiver.callback(broker_message.message)
+
+    result = result_backend.results["task_id"]
+    assert result.is_err
+    assert isinstance(result.error, RuntimeError)
+    assert str(result.error) == "Current message is not ackable."
+    assert events == ["post_execute", "save", "post_save"]
+
+
 async def test_invalid_task_ack_type_raises_before_execution() -> None:
     """Invalid task ack_type fails before task execution and acknowledgement."""
     events: list[str] = []

tests/test_acks.py

@@ -14,6 +14,10 @@ def test_parse_acknowledge_type_from_string() -> None:
     assert parse_acknowledge_type("when_executed") == AcknowledgeType.WHEN_EXECUTED
 
 
+def test_parse_manual_acknowledge_type_from_string() -> None:
+    assert parse_acknowledge_type("manual") == AcknowledgeType.MANUAL
+
+
 def test_parse_acknowledge_type_unknown_string() -> None:
     with pytest.raises(ValueError, match="Unknown acknowledge type value"):
         parse_acknowledge_type("when_save")