Add JetStream acknowledgement progress

README.md

@@ -93,6 +93,30 @@ if __name__ == "__main__":
     asyncio.run(main())
 ```
 
+### JetStream acknowledgement progress
+
+JetStream brokers expose acknowledgement progress for long-running tasks. Call it
+before the consumer's acknowledgement timeout expires to tell JetStream that the
+message is still being processed.
+
+```python
+from taskiq import Context, TaskiqDepends
+
+
+@broker.task(ack_type="manual")
+async def long_running_task(
+    context: Context = TaskiqDepends(),
+) -> None:
+    while still_working:
+        await do_some_work()
+        await context.ack_progress()
+
+    await context.ack()
+```
+
+`ack_progress()` can be called repeatedly. It does not finally acknowledge the
+message; call `ack()` when processing has completed successfully.
+
 ## NatsBroker configuration
 
 Here's the constructor parameters:

taskiq_nats/broker.py

@@ -196,6 +196,7 @@ async def listen(self) -> typing.AsyncGenerator[AckableMessage, None]:
             yield AckableMessage(
                 data=message.data,
                 ack=message.ack,
+                ack_progress=message.in_progress,
             )
 
     async def _startup_consumer(self) -> None:
@@ -237,6 +238,7 @@ async def listen(self) -> typing.AsyncGenerator[AckableMessage, None]:
                     yield AckableMessage(
                         data=nats_message.data,
                         ack=nats_message.ack,
+                        ack_progress=nats_message.in_progress,
                     )
             except NatsTimeoutError:
                 continue

tests/test_jetstream.py

@@ -1,12 +1,54 @@
 import asyncio
 import uuid
+from collections.abc import AsyncGenerator
+from unittest.mock import AsyncMock, MagicMock
 
 from taskiq import AckableMessage, BrokerMessage
 
 from taskiq_nats import PullBasedJetStreamBroker, PushBasedJetStreamBroker
 from tests.utils import read_message
 
 
+async def test_push_based_broker_exposes_ack_progress(
+    nats_urls: list[str],
+) -> None:
+    nats_message = MagicMock()
+    nats_message.data = b"message"
+    nats_message.ack = AsyncMock()
+    nats_message.in_progress = AsyncMock()
+
+    async def message_stream() -> AsyncGenerator[MagicMock, None]:
+        yield nats_message
+
+    broker = PushBasedJetStreamBroker(servers=nats_urls)
+    broker.consumer = MagicMock(messages=message_stream())
+
+    message = await anext(broker.listen())
+
+    assert message.ack_progress is not None
+    await message.ack_progress()
+    nats_message.in_progress.assert_awaited_once_with()
+
+
+async def test_pull_based_broker_exposes_ack_progress(
+    nats_urls: list[str],
+) -> None:
+    nats_message = MagicMock()
+    nats_message.data = b"message"
+    nats_message.ack = AsyncMock()
+    nats_message.in_progress = AsyncMock()
+
+    broker = PullBasedJetStreamBroker(servers=nats_urls)
+    broker.consumer = MagicMock()
+    broker.consumer.fetch = AsyncMock(return_value=[nats_message])
+
+    message = await anext(broker.listen())
+
+    assert message.ack_progress is not None
+    await message.ack_progress()
+    nats_message.in_progress.assert_awaited_once_with()
+
+
 async def test_push_based_broker_success(  # (too many await)
     nats_urls: list[str],
     nats_subject: str,
@@ -36,6 +78,8 @@ async def test_push_based_broker_success(  # (too many await)
     ackable_msg = await asyncio.wait_for(read_message(broker), 0.5)
     assert isinstance(ackable_msg, AckableMessage)
     assert ackable_msg.data == sent_message.message
+    assert ackable_msg.ack_progress is not None
+    await ackable_msg.ack_progress()
     ack = ackable_msg.ack()
     if ack is not None:
         await ack
@@ -76,6 +120,8 @@ async def test_pull_based_broker_success(
     ackable_msg = await asyncio.wait_for(read_message(broker), 0.5)
     assert isinstance(ackable_msg, AckableMessage)
     assert ackable_msg.data == sent_message.message
+    assert ackable_msg.ack_progress is not None
+    await ackable_msg.ack_progress()
     ack = ackable_msg.ack()
     if ack is not None:
         await ack