Run post_send before in-place execution in InMemoryBroker

With InMemoryBroker(await_inplace=True) the middleware hooks fired in the
order pre_send -> pre_execute -> task -> post_execute -> post_send, because
kick() awaited the full receiver callback (pre_execute/task/post_execute)
before control returned to the sender that invokes post_send.

Now kick() schedules the in-place receiver as a task and returns
immediately, and the task is awaited in a new _finish_kick() hook that runs
right after post_send. This restores the expected ordering
pre_send -> post_send -> pre_execute -> task -> post_execute while the task
is still fully executed before kiq() returns.

Fixes #586.

Author: apoorvdarshan

taskiq/abc/broker.py

@@ -237,6 +237,18 @@ async def kick(
         :param message: name of a task.
         """
 
+    async def _finish_kick(self) -> None:  # noqa: B027
+        """
+        Hook that runs after a message is sent and ``post_send`` fired.
+
+        Most brokers send messages elsewhere and have nothing to do here,
+        so this is a no-op by default. Brokers that execute tasks in-place
+        (like :class:`~taskiq.InMemoryBroker` with ``await_inplace``) can
+        override this method to await the in-place execution *after* the
+        ``post_send`` middleware hook has run, preserving the expected
+        ``pre_send -> post_send -> pre_execute -> ...`` ordering.
+        """
+
     @abstractmethod
     def listen(self) -> AsyncGenerator[bytes | AckableMessage, None]:
         """

taskiq/brokers/inmemory_broker.py

@@ -146,6 +146,7 @@ def __init__(
         )
         self.await_inplace = await_inplace
         self._running_tasks: set[asyncio.Task[Any]] = set()
+        self._inplace_tasks: set[asyncio.Task[Any]] = set()
 
     async def kick(self, message: BrokerMessage) -> None:
         """
@@ -162,14 +163,37 @@ async def kick(self, message: BrokerMessage) -> None:
             raise UnknownTaskError(task_name=message.task_name)
 
         receiver_cb = self.receiver.callback(message=message.message)
+        task = asyncio.create_task(receiver_cb)
+
         if self.await_inplace:
-            await receiver_cb
+            # Schedule the receiver as a task instead of awaiting it here so
+            # that control returns to the sender and the ``post_send``
+            # middleware hook fires *before* the task starts executing. The
+            # task is then awaited in `_finish_kick`, right after `post_send`,
+            # keeping the in-place execution while fixing the hook ordering.
+            self._inplace_tasks.add(task)
+            task.add_done_callback(self._inplace_tasks.discard)
             return
 
-        task = asyncio.create_task(receiver_cb)
         self._running_tasks.add(task)
         task.add_done_callback(self._running_tasks.discard)
 
+    async def _finish_kick(self) -> None:
+        """
+        Await in-place task execution scheduled by `kick`.
+
+        This runs after the `post_send` middleware hook, so with
+        ``await_inplace=True`` the hook ordering becomes
+        ``pre_send -> post_send -> pre_execute -> task -> post_execute``
+        while tasks are still fully executed before `kiq` returns.
+        """
+        if not self._inplace_tasks:
+            return
+        to_await = list(self._inplace_tasks)
+        self._inplace_tasks.clear()
+        for task in to_await:
+            await task
+
     def listen(self) -> AsyncGenerator[bytes, None]:
         """
         Inmemory broker cannot listen.

taskiq/kicker.py

@@ -169,6 +169,8 @@ async def kiq(
             if middleware.__class__.post_send != TaskiqMiddleware.post_send:
                 await maybe_awaitable(middleware.post_send(message))
 
+        await self.broker._finish_kick()  # noqa: SLF001
+
         return AsyncTaskiqTask(
             task_id=message.task_id,
             result_backend=self.broker.result_backend,

tests/brokers/test_inmemory.py

@@ -4,7 +4,10 @@
 import pytest
 
 from taskiq import InMemoryBroker
+from taskiq.abc.middleware import TaskiqMiddleware
 from taskiq.events import TaskiqEvents
+from taskiq.message import TaskiqMessage
+from taskiq.result import TaskiqResult
 from taskiq.state import TaskiqState
 
 
@@ -114,3 +117,46 @@ async def test_task() -> None:
     assert slept
     assert await task.is_ready()
     assert not broker._running_tasks
+
+
+async def test_inplace_middleware_order() -> None:
+    """post_send must fire before task execution with await_inplace=True."""
+    events: list[str] = []
+
+    class OrderMiddleware(TaskiqMiddleware):
+        async def pre_send(self, message: TaskiqMessage) -> TaskiqMessage:
+            events.append("pre_send")
+            return message
+
+        async def post_send(self, message: TaskiqMessage) -> None:
+            events.append("post_send")
+
+        async def pre_execute(self, message: TaskiqMessage) -> TaskiqMessage:
+            events.append("pre_execute")
+            return message
+
+        async def post_execute(
+            self,
+            message: TaskiqMessage,
+            result: "TaskiqResult[object]",
+        ) -> None:
+            events.append("post_execute")
+
+    broker = InMemoryBroker(await_inplace=True).with_middlewares(OrderMiddleware())
+
+    @broker.task
+    async def test_task() -> None:
+        events.append("task")
+
+    task = await test_task.kiq()
+
+    # The task must be fully executed by the time kiq returns.
+    assert await task.is_ready()
+    assert not broker._inplace_tasks
+    assert events == [
+        "pre_send",
+        "post_send",
+        "pre_execute",
+        "task",
+        "post_execute",
+    ]