fix: close the root container after all worker shutdown code

CONTEXT.md

@@ -18,8 +18,9 @@ package's own.
 
 **Root container**:
 The `Container` that `setup_di` attaches to `broker.state`: the one a caller constructs and hands
-in, and the one `fetch_di_container` returns. By default the worker's startup and shutdown events
-open and close it; with `setup_di(..., manage_lifespan=False)` the caller does.
+in, and the one `fetch_di_container` returns. By default the worker's startup event opens it, and
+shutdown closes it after all `WORKER_SHUTDOWN` handlers and middleware `shutdown()` calls; with
+`setup_di(..., manage_lifespan=False)` the caller does.
 _Avoid_: APP-scope container — that names a scope where what matters is the position at the top of
 the tree, and `Scope.APP` is `modern-di`'s word for a band in the hierarchy, not for this object.
 

README.md

@@ -31,7 +31,7 @@ uv add modern-di-taskiq      # or: pip install modern-di-taskiq
 
 ## Usage
 
-`setup_di` stores the container on `broker.state` and registers `WORKER_STARTUP`/`WORKER_SHUTDOWN` handlers that open/close it. `FromDI` resolves a provider (or type) into a task parameter through a taskiq dependency, which builds a `Scope.REQUEST` child container for the task and closes it when the task finishes. No per-task decorator is needed.
+`setup_di` stores the container on `broker.state` and opens it on `WORKER_STARTUP`. At shutdown it closes the container after every `WORKER_SHUTDOWN` handler and every middleware `shutdown()`, so your shutdown code can still resolve `Scope.APP` dependencies and their finalizers run last. `FromDI` resolves a provider (or type) into a task parameter through a taskiq dependency, which builds a `Scope.REQUEST` child container for the task and closes it when the task finishes. No per-task decorator is needed.
 
 ```python
 import typing

modern_di_taskiq/main.py

@@ -3,7 +3,7 @@
 
 import taskiq
 from modern_di import Container, Scope, integrations, providers
-from taskiq import AsyncBroker, Context, TaskiqDepends, TaskiqEvents, TaskiqState
+from taskiq import AsyncBroker, Context, TaskiqDepends, TaskiqEvents, TaskiqMiddleware, TaskiqState
 
 
 taskiq_message_provider = providers.ContextProvider(taskiq.TaskiqMessage, scope=Scope.REQUEST)
@@ -16,17 +16,46 @@
 _CONNECTION_PROVIDERS = (taskiq_message_provider,)
 
 
+class _RootCloser(TaskiqMiddleware):
+    def __init__(self, container: Container) -> None:
+        super().__init__()
+        self.container = container
+        self.broker_runs_middlewares = False
+
+    def startup(self) -> None:
+        self.broker_runs_middlewares = True
+
+    async def shutdown(self) -> None:
+        if self.broker_runs_middlewares:
+            await self.container.close_async()
+
+
+def _move_to_end(items: list[typing.Any], item: object) -> None:
+    if item in items:
+        items.remove(item)
+    items.append(item)
+
+
 def setup_di(broker: AsyncBroker, container: Container, *, manage_lifespan: bool = True) -> Container:
     setattr(broker.state, _ROOT_CONTAINER_ATTR, container)
     container.add_providers(*_CONNECTION_PROVIDERS)
     if not manage_lifespan:
         return container
 
+    closer = _RootCloser(container)
+
     def _open_root(_state: TaskiqState) -> None:
         container.open()
+        closer.broker_runs_middlewares = False
+        if closer not in broker.middlewares:
+            broker.add_middlewares(closer)
+        _move_to_end(broker.event_handlers[TaskiqEvents.WORKER_SHUTDOWN], _close_root)
 
     async def _close_root(_state: TaskiqState) -> None:
-        await container.close_async()
+        if closer.broker_runs_middlewares and closer in broker.middlewares:
+            _move_to_end(broker.middlewares, closer)
+        else:
+            await container.close_async()
 
     broker.add_event_handler(TaskiqEvents.WORKER_STARTUP, _open_root)
     broker.add_event_handler(TaskiqEvents.WORKER_SHUTDOWN, _close_root)

pyproject.toml

@@ -17,7 +17,7 @@ classifiers = [
     "Typing :: Typed",
     "Topic :: Software Development :: Libraries",
 ]
-dependencies = ["taskiq>=0.11,<0.13", "modern-di>=4,<5"]
+dependencies = ["taskiq>=0.11.8,<0.13", "modern-di>=4,<5"]
 version = "0"
 
 [project.urls]

tests/test_shutdown_order.py

@@ -0,0 +1,148 @@
+import collections.abc
+
+import pytest
+from modern_di import Container, Group, Scope, providers
+from taskiq import AsyncBroker, BrokerMessage, InMemoryBroker, TaskiqEvents, TaskiqMiddleware, TaskiqState
+
+from modern_di_taskiq import setup_di
+
+
+class Pool:
+    pass
+
+
+class StubBroker(AsyncBroker):
+    def __init__(self) -> None:
+        super().__init__()
+        self.is_worker_process = True
+
+    async def kick(self, message: BrokerMessage) -> None:  # pragma: no cover
+        pass
+
+    async def listen(self) -> collections.abc.AsyncGenerator[bytes]:  # pragma: no cover
+        yield b""
+
+
+def _container(events: list[str]) -> Container:
+    class Resources(Group):
+        pool = providers.Factory(
+            Pool, scope=Scope.APP, cache=providers.CacheSettings(finalizer=lambda _: events.append("finalized"))
+        )
+
+    return Container(groups=[Resources])
+
+
+def _use(container: Container, events: list[str], label: str) -> None:
+    container.resolve(Pool)
+    events.append(label)
+
+
+class UserMiddleware(TaskiqMiddleware):
+    def __init__(self, container: Container, events: list[str], label: str) -> None:
+        super().__init__()
+        self.container = container
+        self.events = events
+        self.label = label
+
+    async def shutdown(self) -> None:
+        _use(self.container, self.events, self.label)
+
+
+def _register_user_code(broker: AsyncBroker, container: Container, events: list[str], label: str) -> None:
+    async def hook(_state: TaskiqState) -> None:
+        _use(container, events, f"{label} hook")
+
+    broker.add_event_handler(TaskiqEvents.WORKER_SHUTDOWN, hook)
+    broker.add_middlewares(UserMiddleware(container, events, f"{label} middleware"))
+
+
+async def test_root_closes_after_all_shutdown_handlers_and_middlewares() -> None:
+    events: list[str] = []
+    broker = StubBroker()
+    container = _container(events)
+    _register_user_code(broker, container, events, "before")
+    setup_di(broker, container)
+    _register_user_code(broker, container, events, "after")
+
+    await broker.startup()
+    container.resolve(Pool)
+    await broker.shutdown()
+
+    assert events == ["before hook", "after hook", "before middleware", "after middleware", "finalized"]
+    assert container.closed is True
+
+
+async def test_root_closes_after_a_middleware_added_after_startup() -> None:
+    events: list[str] = []
+    broker = StubBroker()
+    container = setup_di(broker, _container(events))
+
+    await broker.startup()
+    broker.add_middlewares(UserMiddleware(container, events, "late middleware"))
+    container.resolve(Pool)
+    await broker.shutdown()
+
+    assert events == ["late middleware", "finalized"]
+
+
+async def test_root_closes_once_per_cycle_without_duplicate_middlewares() -> None:
+    events: list[str] = []
+    broker = StubBroker()
+    container = setup_di(broker, _container(events))
+    _register_user_code(broker, container, events, "user")
+
+    for _ in range(2):
+        await broker.startup()
+        middleware_count = len(broker.middlewares)
+        container.resolve(Pool)
+        await broker.shutdown()
+        assert len(broker.middlewares) == middleware_count
+        assert container.closed is True
+
+    assert events == ["user hook", "user middleware", "finalized"] * 2
+
+
+async def test_root_closes_once_when_the_broker_skips_middleware_shutdown() -> None:
+    events: list[str] = []
+    broker = InMemoryBroker()
+    container = setup_di(broker, _container(events))
+
+    async def hook(_state: TaskiqState) -> None:
+        _use(container, events, "hook")
+
+    broker.add_event_handler(TaskiqEvents.WORKER_SHUTDOWN, hook)
+
+    for _ in range(2):
+        await broker.startup()
+        container.resolve(Pool)
+        await broker.shutdown()
+        assert container.closed is True
+
+    assert events == ["hook", "finalized"] * 2
+
+
+async def test_root_closes_when_shutdown_runs_without_startup() -> None:
+    events: list[str] = []
+    broker = StubBroker()
+    container = setup_di(broker, _container(events))
+    container.resolve(Pool)
+
+    await broker.shutdown()
+
+    assert container.closed is True
+    assert events == ["finalized"]
+
+
+@pytest.mark.parametrize("broker_type", [StubBroker, InMemoryBroker])
+async def test_manage_lifespan_false_adds_no_middleware_and_keeps_the_root_open(broker_type: type[AsyncBroker]) -> None:
+    events: list[str] = []
+    broker = broker_type()
+    container = setup_di(broker, _container(events), manage_lifespan=False)
+    container.resolve(Pool)
+
+    await broker.startup()
+    assert broker.middlewares == []
+    await broker.shutdown()
+
+    assert container.closed is False
+    assert events == []

tests/test_taskiq_di.py

@@ -10,7 +10,7 @@
 
 async def test_resolves_app_request_and_context(broker: InMemoryBroker) -> None:
     # Explicit startup/shutdown (not `async with broker`) so the suite runs on
-    # the whole taskiq>=0.11 floor — InMemoryBroker gained async-context-manager
+    # the whole taskiq>=0.11.8 floor — InMemoryBroker gained async-context-manager
     # support only in 0.12.3, and the integration itself needs neither.
     @broker.task(task_name="my_task")
     async def my_task(