From 57467b3c7e43d3a7edd4b580c5ef65799bc39ca4 Mon Sep 17 00:00:00 2001 From: Artur Shiriev Date: Tue, 15 Sep 2026 19:57:04 +0300 Subject: [PATCH] fix: install the DI middleware on every broker, on startup FastStream 0.7 apps hold a list of brokers and app.broker is only the first, so setup_di left brokers 2..N without DI. The first message to a FromDI subscriber on one of them failed with a None request container and nothing at setup time pointing at the cause. setup_di now registers an on_startup hook that adds the middleware to each broker in app.brokers not yet carrying it. Startup is when the broker list is complete, so a broker added via app.add_broker after setup_di is covered too, and the membership check keeps a restarted app from stacking a second copy. Rationale in docs/adr/0002-install-middleware-on-startup.md. Closes #42 --- README.md | 2 +- .../adr/0002-install-middleware-on-startup.md | 40 +++++++++++++ modern_di_faststream/main.py | 11 +++- tests/test_faststream_di.py | 56 +++++++++++++++++++ 4 files changed, 107 insertions(+), 2 deletions(-) create mode 100644 docs/adr/0002-install-middleware-on-startup.md diff --git a/README.md b/README.md index a121830..22e9d21 100644 --- a/README.md +++ b/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` diff --git a/docs/adr/0002-install-middleware-on-startup.md b/docs/adr/0002-install-middleware-on-startup.md new file mode 100644 index 0000000..7340cdc --- /dev/null +++ b/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. diff --git a/modern_di_faststream/main.py b/modern_di_faststream/main.py index 93b85d2..cab0863 100644 --- a/modern_di_faststream/main.py +++ b/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 diff --git a/tests/test_faststream_di.py b/tests/test_faststream_di.py index d0b5bb3..1de1797 100644 --- a/tests/test_faststream_di.py +++ b/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