fix: install the DI middleware on every broker, on startup

README.md

@@ -73,7 +73,7 @@ The current `StreamMessage` is resolvable within DI via the pre-built `faststrea
 
 ## API
 
-- `setup_di(app, container)` — stores the container in the app context, registers startup/shutdown lifecycle hooks (reopen on startup, close after shutdown), and adds the DI middleware to the broker
+- `setup_di(app, container)` — 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`), after shutdown it closes the container
 - `FromDI(dependency, *, use_cache=True, cast=False)` — FastStream `Depends` that resolves a provider (or type) from the request container
 - `fetch_di_container(app)` — returns the root container from the app context
 - `faststream_message_provider` — `ContextProvider` for the current `faststream.StreamMessage`

docs/adr/0002-install-middleware-on-startup.md

@@ -0,0 +1,40 @@
+# Install the DI middleware on startup, on every broker
+
+**Decision:** `setup_di` does not call `add_middleware` itself. It registers an `on_startup` hook
+that walks `app.brokers` and adds the middleware factory to each broker that does not already
+carry it.
+
+## Why
+
+FastStream 0.7 apps hold a list of brokers. `FastStream(*brokers)` accepts many, `app.add_broker`
+appends more after construction, and `app.broker` is only `brokers[0]`. Installing on `app.broker`
+at `setup_di` time therefore left every other broker without DI, and the gap was silent: the app
+started, and the first message to a subscriber on another broker failed inside `FromDI` with a
+missing request container and nothing pointing at the cause
+([#42](https://github.com/modern-python/modern-di-faststream/issues/42)).
+
+Iterating `app.brokers` inside `setup_di` fixes the construction-time case but still misses a
+broker added afterwards, and the only remedy would be a documented ordering rule the user has to
+remember. Startup is the one moment when the broker list is complete and no message has been
+consumed yet, so installing there needs no rule. It is safe because FastStream builds a
+subscriber's middleware stack per message from the broker config, so a middleware added in an
+`on_startup` hook applies to subscribers registered before it.
+
+The membership check exists because `on_startup` runs on every start. Without it a stopped and
+restarted app would carry two copies and build two request containers per message. The check
+reads `broker.config.broker_middlewares`, the same sequence FastStream itself builds the stack
+from, rather than a private record of installed brokers that could drift from it.
+
+## What changes for a reader
+
+Between `setup_di` and startup the middleware is not yet on any broker. Nothing in this package,
+its tests, or its documentation inspects a broker in that window; `TestApp` runs the startup hooks.
+
+The `if not app.broker` guard in `setup_di` is kept as it was. Relaxing it so that a broker
+created inside the user's own `on_startup` hook can be picked up is possible, since hooks run in
+registration order, but that ordering is subtle enough to want its own decision.
+
+**Revisit trigger:** FastStream exposes a hook for a broker being added to an app, so the
+middleware can be installed at that moment instead of on startup, **or** FastStream freezes a
+subscriber's middleware stack before `on_startup` runs, which would make a startup-time install
+too late.

modern_di_faststream/main.py

@@ -65,13 +65,22 @@ def setup_di(
 
     container.add_providers(faststream_message_provider)
     app.context.set_global(_ROOT_CONTAINER_KEY, container)
+    middleware_factory = _DIMiddlewareFactory(container)
+
+    # Installed on startup rather than here so a broker added after ``setup_di`` is covered,
+    # and skipped where present so a restart adds no second copy: docs/adr/0002-install-middleware-on-startup.md
+    def install_middleware() -> None:
+        for broker in app.brokers:
+            if middleware_factory not in broker.config.broker_middlewares:
+                broker.add_middleware(middleware_factory)
+
     # 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.
     app.on_startup(container.open)
+    app.on_startup(install_middleware)
     app.after_shutdown(container.close_async)
-    app.broker.add_middleware(_DIMiddlewareFactory(container))
     return container
 
 

tests/test_faststream_di.py

@@ -8,6 +8,7 @@
 
 import modern_di_faststream
 from modern_di_faststream import FromDI
+from modern_di_faststream.main import _DIMiddlewareFactory
 from tests.dependencies import Dependencies, DependentCreator, SimpleCreator
 
 
@@ -55,3 +56,58 @@ async def test_app_without_broker() -> None:
 def test_fetch_di_container(app: faststream.FastStream) -> None:
     di_container = modern_di_faststream.fetch_di_container(app)
     assert isinstance(di_container, Container)
+
+
+def _app_with_two_brokers(*, add_second_after_setup: bool) -> tuple[faststream.FastStream, NatsBroker, NatsBroker]:
+    first, second = NatsBroker(), NatsBroker()
+    app_ = faststream.FastStream(first) if add_second_after_setup else faststream.FastStream(first, second)
+    modern_di_faststream.setup_di(app_, container=Container(groups=[Dependencies]))
+    if add_second_after_setup:
+        app_.add_broker(second)
+    return app_, first, second
+
+
+def _subscribe_resolving(broker: NatsBroker, subject: str, resolved: list[SimpleCreator]) -> None:
+    @broker.subscriber(subject)
+    async def subscriber(instance: typing.Annotated[SimpleCreator, FromDI(Dependencies.app_factory)]) -> None:
+        resolved.append(instance)
+
+
+@pytest.mark.parametrize("add_second_after_setup", [False, True], ids=["at-construction", "after-setup_di"])
+async def test_di_resolves_on_every_broker(add_second_after_setup: bool) -> None:
+    """INVARIANT: a ``FromDI`` parameter resolves on every broker of the app, not only ``app.broker``.
+
+    Broken by installing the middleware on ``app.broker`` alone, or by installing it when
+    ``setup_di`` runs instead of on startup: FastStream 0.7 apps hold a list of brokers, and
+    ``app.add_broker`` may append to it after ``setup_di`` returned. Either regression is silent at
+    setup and surfaces as a missing request container on the first message to the other broker.
+    """
+    app_, first, second = _app_with_two_brokers(add_second_after_setup=add_second_after_setup)
+    resolved: list[SimpleCreator] = []
+    _subscribe_resolving(first, "first", resolved)
+    _subscribe_resolving(second, "second", resolved)
+
+    async with TestNatsBroker(first, second) as (first_test, second_test), TestApp(app_):
+        await first_test.publish(None, "first")
+        await second_test.publish(None, "second")
+
+    assert [type(instance) for instance in resolved] == [SimpleCreator, SimpleCreator]
+
+
+async def test_middleware_is_installed_once_per_broker_across_restarts() -> None:
+    """INVARIANT: each broker carries exactly one DI middleware however many times the app starts.
+
+    Broken by a startup hook that adds the middleware unconditionally. Every start after the first
+    would then add another copy, and each message would build one request container per copy,
+    with only the innermost visible to ``FromDI``.
+    """
+    app_, first, second = _app_with_two_brokers(add_second_after_setup=False)
+
+    async with TestNatsBroker(first, second), TestApp(app_):
+        pass
+    async with TestNatsBroker(first, second), TestApp(app_):
+        pass
+
+    for broker in (first, second):
+        installed = [m for m in broker.config.broker_middlewares if isinstance(m, _DIMiddlewareFactory)]
+        assert len(installed) == 1