fix: close the root container after every after_shutdown hook (#83)

Author: lesnik512

README.md

@@ -73,7 +73,7 @@ The current `StreamMessage` is resolvable within DI via the pre-built `faststrea
 
 ### Sharing a container
 
-By default the app opens the container on startup and closes it after shutdown, which runs every `APP`-scoped finalizer. When the container outlives the app, or other apps share it, pass `manage_lifespan=False` and open and close it yourself:
+By default the app opens the container on startup and closes it after shutdown, which runs every `APP`-scoped finalizer. The close runs after every `on_shutdown` and `after_shutdown` hook, including hooks registered after `setup_di`, so they can still resolve `APP`-scoped dependencies. The teardown of a `FastStream(lifespan=...)` context manager runs after the close, because FastStream exits it after all hooks. If that teardown needs the container, or the container outlives the app, or other apps share it, pass `manage_lifespan=False` and open and close it yourself:
 
 ```python
 container = Container(groups=[Dependencies])
@@ -114,7 +114,7 @@ FastStream's FastAPI routers, such as `faststream.nats.fastapi.NatsRouter`, need
 
 | Symbol | Description |
 |---|---|
-| `setup_di(app, container, *, manage_lifespan=True)` | Stores the container in the app context and registers startup/shutdown lifecycle hooks. On startup it reopens the container and adds the DI middleware to every broker of the app, including one added via `app.add_broker` after `setup_di` or by an `on_startup` hook registered before `setup_di`. After shutdown it closes the container. With `manage_lifespan=False` it never opens or closes the container, and the middleware still builds and closes a request container per message. Raises `RuntimeError` at startup if the app has no broker by then. Returns the container |
+| `setup_di(app, container, *, manage_lifespan=True)` | Stores the container in the app context and registers startup/shutdown lifecycle hooks. On startup it reopens the container and adds the DI middleware to every broker of the app, including one added via `app.add_broker` after `setup_di` or by an `on_startup` hook registered before `setup_di`. After the `after_shutdown` hooks run, it closes the container. With `manage_lifespan=False` it never opens or closes the container, and the middleware still builds and closes a request container per message. Raises `RuntimeError` at startup if the app has no broker by then. Returns the container |
 | `FromDI(dependency, *, use_cache=True, cast=False)` | FastStream `Depends` that resolves a provider (or type) from the request container. Raises `RuntimeError` naming `setup_di` when a message reaches it without the middleware installed |
 | `fetch_di_container(app)` | Returns the root container from the app context |
 | `faststream_message_provider` | `ContextProvider` for the current `faststream.StreamMessage` (`REQUEST` scope) |

modern_di_faststream/main.py

@@ -87,13 +87,18 @@ def install_middleware() -> None:
                 broker.add_middleware(middleware_factory)
 
     app.on_startup(install_middleware)
-    # FastStream's lifecycle is callback-based, so the root container can't be
-    # wrapped in ``async with``. Reopen it on startup (before the broker consumes)
-    # to pair with the shutdown close, so a broker restart works instead of
-    # raising ContainerClosedError. Reopening an already-open container is a no-op.
     if manage_lifespan:
-        app.on_startup(container.open)
-        app.after_shutdown(container.close_async)
+        close_registered = False
+
+        # Appended on startup so the close runs after every after_shutdown hook registered before it.
+        def open_root() -> None:
+            nonlocal close_registered
+            container.open()
+            if not close_registered:
+                app.after_shutdown(container.close_async)
+                close_registered = True
+
+        app.on_startup(open_root)
     return container
 
 

tests/test_lifespan.py

@@ -1,7 +1,10 @@
 import typing
+from collections.abc import Callable
 
 import faststream
+import pytest
 from faststream import TestApp
+from faststream.asgi import AsgiFastStream
 from faststream.nats import NatsBroker, TestNatsBroker
 from modern_di import Container
 
@@ -95,3 +98,51 @@ async def test_manage_lifespan_false_lets_the_user_own_the_root() -> None:
 
     assert container.closed
     assert app_instance.finalized == 1
+
+
+class _CloseCountingContainer(Container):
+    close_calls = 0
+
+    async def close_async(self) -> None:
+        self.close_calls += 1
+        await super().close_async()
+
+
+def _setup_on_faststream(broker: NatsBroker, container: Container) -> faststream.FastStream | AsgiFastStream:
+    app = faststream.FastStream(broker)
+    setup_di(app, container)
+    return app
+
+
+def _setup_on_asgi(broker: NatsBroker, container: Container) -> faststream.FastStream | AsgiFastStream:
+    app = AsgiFastStream(broker)
+    setup_di(app, container)
+    return app
+
+
+def _setup_then_as_asgi(broker: NatsBroker, container: Container) -> faststream.FastStream | AsgiFastStream:
+    app = faststream.FastStream(broker)
+    setup_di(app, container)
+    return app.as_asgi()
+
+
+@pytest.mark.parametrize("build_app", [_setup_on_faststream, _setup_on_asgi, _setup_then_as_asgi])
+async def test_after_shutdown_hook_registered_after_setup_di_sees_the_open_root(
+    build_app: Callable[[NatsBroker, Container], faststream.FastStream | AsgiFastStream],
+) -> None:
+    broker = NatsBroker()
+    container = _CloseCountingContainer(groups=[Dependencies])
+    app = build_app(broker, container)
+    seen: list[tuple[Finalizable, int]] = []
+
+    @app.after_shutdown
+    def user_hook() -> None:
+        instance = container.resolve_provider(Dependencies.app_finalizable)
+        seen.append((instance, instance.finalized))
+
+    for cycle in (1, 2):
+        async with TestNatsBroker(broker), TestApp(app):
+            pass
+        assert [(instance.finalized, at_hook) for instance, at_hook in seen] == [(1, 0)] * cycle
+        assert container.closed
+        assert container.close_calls == cycle