fix(inmemory): stabilize task ordering and lifecycle

docs/available-components/brokers.md

@@ -8,10 +8,21 @@ In this section we'll list officially supported brokers.
 
 ## InMemoryBroker
 
-This is a special broker for local development. It uses the same functions to execute tasks,
-but all tasks are executed locally in the current thread.
+This is a special broker for local development. It uses the same functions to
+execute tasks in the current process. Async task functions run on the event
+loop, while sync task functions use the broker's thread pool.
 By default it uses `InMemoryResultBackend` but this can be overridden.
 
+`startup()` and `shutdown()` run both the client and worker event phases because
+the broker performs both roles in one process. They also own middleware and
+result backend lifecycle. Shutdown stops accepting new local executions, waits
+for already accepted work (including active send middleware), then closes those
+resources and the synchronous task executor. A send started after shutdown is
+rejected before its `pre_send` hooks run.
+Once cleanup completes, later `shutdown()` calls do not close those resources
+again. If cancellation interrupts a lifecycle hook, it propagates immediately
+and a later call resumes cleanup from that hook.
+
 ## ZeroMQBroker
 
 This broker uses [ZMQ](https://zeromq.org/) to communicate between worker and client processes.

docs/guide/testing-taskiq.md

@@ -165,6 +165,13 @@ broker = InMemoryBroker(await_inplace=True)
 
 With this setup all `await function.kiq()` calls will behave similarly to `await function()`, but
 with dependency injection and all taskiq-related functionality.
+For async task functions, successful inline execution remains in the caller's
+asyncio task and context. Sync task functions still run in the broker's thread
+pool with a copy of the current `ContextVars`.
+The normal middleware boundary is preserved in both execution modes: all
+`post_send` hooks finish before `pre_execute` and the task body begin.
+With `await_inplace=True`, the `kiq()` call also returns only after inline
+execution completes.
 
 2. Alternatively, you can manually await all tasks after invoking the
    target function by using the `wait_all` method.
@@ -182,6 +189,25 @@ async def test_add_one():
     # have been completed and do all the assertions.
 ```
 
+`wait_all()` also surfaces the first pending background callback failure even
+if that callback completed before `wait_all()` was called. The failure is
+consumed after it is raised, so a later `wait_all()` only observes newer work.
+Cancelling a `wait_all()` call cancels only that waiter; accepted executions
+remain tracked and can be drained by a later call.
+Calling `shutdown()` rejects new sends, performs the same drain, then closes
+middleware, result backend, and executor resources. If shutdown is cancelled
+while draining accepted work, it finishes that drain and cleanup before
+propagating the cancellation.
+Cancellation from a shutdown event, middleware, or result backend hook instead
+propagates immediately. A later `shutdown()` resumes at the interrupted hook
+without repeating completed cleanup, while calls after completed cleanup are
+no-ops.
+An invocation already running `pre_send` is part of that drain; a later
+invocation is rejected before `pre_send` can produce side effects.
+Both drain methods must be called by the external test or application lifecycle
+owner. Calling `wait_all()` or `shutdown()` from a task or `post_send` hook
+managed by the same broker raises `RuntimeError` instead of waiting on itself.
+
 ## Dependency injection
 
 If you use dependencies in your tasks, you may think that this can become a problem for tests. But it's not.

taskiq/abc/broker.py

@@ -3,8 +3,9 @@
 import warnings
 from abc import ABC, abstractmethod
 from collections import defaultdict
-from collections.abc import AsyncGenerator, Awaitable, Callable
-from functools import wraps
+from collections.abc import AsyncGenerator, Awaitable, Callable, Iterator
+from contextlib import contextmanager
+from functools import partial, wraps
 from logging import getLogger
 from typing import (
     TYPE_CHECKING,
@@ -23,7 +24,7 @@
 from taskiq.acks import AckableMessage
 from taskiq.decor import AsyncTaskiqDecoratedTask
 from taskiq.events import TaskiqEvents
-from taskiq.exceptions import TaskBrokerMismatchError
+from taskiq.exceptions import SendTaskError, TaskBrokerMismatchError
 from taskiq.formatters.proxy_formatter import ProxyFormatter
 from taskiq.message import BrokerMessage
 from taskiq.result_backends.dummy import DummyResultBackend
@@ -47,6 +48,7 @@
 _ReturnType = TypeVar("_ReturnType")
 
 EventHandler: TypeAlias = Callable[[TaskiqState], Awaitable[None] | None]
+ShutdownHook: TypeAlias = Callable[[], Awaitable[None] | None]
 
 logger = getLogger("taskiq")
 
@@ -118,6 +120,7 @@ def __init__(
         self.is_worker_process = False
         # True only if broker runs in scheduler process.
         self.is_scheduler_process = False
+        self._shutdown_resource_index = 0
 
     def find_task(self, task_name: str) -> AsyncTaskiqDecoratedTask[Any, Any] | None:
         """
@@ -186,39 +189,109 @@ def add_middlewares(self, *middlewares: "TaskiqMiddleware") -> None:
 
     async def startup(self) -> None:
         """Do something when starting broker."""
-        event = TaskiqEvents.CLIENT_STARTUP
-        if self.is_worker_process:
-            event = TaskiqEvents.WORKER_STARTUP
-
-        for handler in self.event_handlers[event]:
-            await maybe_awaitable(handler(self.state))
+        for event in self._get_startup_events():
+            for handler in self.event_handlers[event]:
+                await maybe_awaitable(handler(self.state))
 
         for middleware in self.middlewares:
             if middleware.__class__.startup != TaskiqMiddleware.startup:
                 await maybe_awaitable(middleware.startup())
 
         await self.result_backend.startup()
 
+    def _get_startup_events(self) -> tuple[TaskiqEvents, ...]:
+        """Return event phases owned by this broker startup."""
+        if self.is_worker_process:
+            return (TaskiqEvents.WORKER_STARTUP,)
+        return (TaskiqEvents.CLIENT_STARTUP,)
+
     async def shutdown(self) -> None:
         """
         Close the broker.
 
         This method is called,
         when broker is closing.
         """
-        event = TaskiqEvents.CLIENT_SHUTDOWN
-        if self.is_worker_process:
-            event = TaskiqEvents.WORKER_SHUTDOWN
+        shutdown_errors: list[BaseException] = []
+        await self._shutdown_resources(shutdown_errors)
 
-        # Call all shutdown events.
-        for handler in self.event_handlers[event]:
-            await maybe_awaitable(handler(self.state))
+        if shutdown_errors:
+            raise shutdown_errors[0]
+
+    async def _shutdown_resources(
+        self,
+        shutdown_errors: list[BaseException],
+    ) -> None:
+        """Close every registered resource and record failures in order."""
+        for index, shutdown_hook in enumerate(self._iter_shutdown_hooks()):
+            if index < self._shutdown_resource_index:
+                continue
+            try:
+                await maybe_awaitable(shutdown_hook())
+            except Exception as exc:
+                self._record_shutdown_error(shutdown_errors, exc)
+            self._shutdown_resource_index = index + 1
+
+    def _iter_shutdown_hooks(self) -> Iterator[ShutdownHook]:
+        """Yield lifecycle hooks in their shutdown order."""
+        for event in self._get_shutdown_events():
+            for handler in self.event_handlers[event]:
+                yield partial(handler, self.state)
 
         for middleware in self.middlewares:
             if middleware.__class__.shutdown != TaskiqMiddleware.shutdown:
-                await maybe_awaitable(middleware.shutdown())
+                yield middleware.shutdown
 
-        await self.result_backend.shutdown()
+        yield self.result_backend.shutdown
+
+    def _get_shutdown_events(self) -> tuple[TaskiqEvents, ...]:
+        """Return event phases owned by this broker shutdown."""
+        if self.is_worker_process:
+            return (TaskiqEvents.WORKER_SHUTDOWN,)
+        return (TaskiqEvents.CLIENT_SHUTDOWN,)
+
+    @staticmethod
+    def _remember_shutdown_error(
+        first_error: BaseException | None,
+        current_error: BaseException,
+    ) -> BaseException:
+        """Keep the first shutdown failure while cleanup continues."""
+        if first_error is None:
+            return current_error
+        logger.error(
+            "Additional error while shutting down broker resources.",
+            exc_info=current_error,
+        )
+        return first_error
+
+    @classmethod
+    def _record_shutdown_error(
+        cls,
+        shutdown_errors: list[BaseException],
+        current_error: BaseException,
+    ) -> None:
+        """Record the first failure when it occurs and log later failures."""
+        first_error = shutdown_errors[0] if shutdown_errors else None
+        remembered_error = cls._remember_shutdown_error(first_error, current_error)
+        if first_error is None:
+            shutdown_errors.append(remembered_error)
+
+    @contextmanager
+    def _send_lifecycle(self) -> Iterator[None]:
+        """Own package-internal client send work through broker handoff."""
+        yield
+
+    async def _kick_with_post_send(
+        self,
+        message: BrokerMessage,
+        post_send: Callable[[], Awaitable[None]],
+    ) -> None:
+        """Run the package-internal send boundary used by AsyncKicker."""
+        try:
+            await self.kick(message)
+        except Exception as exc:
+            raise SendTaskError from exc
+        await post_send()
 
     @abstractmethod
     async def kick(