fix: keep the request container open until a streamed response closes

CONTEXT.md

@@ -29,8 +29,9 @@ Flask app context.
 **Child container**:
 The `Scope.REQUEST` container built per request, on first use or in `before_request`, whichever
 comes first, and stashed on `flask.g`; `inject` resolves from it and `teardown_appcontext` closes
-it after every other `teardown_appcontext` handler. Its lifetime is bounded by Flask's
-**app context**, not by Flask's request context: an app context pushed without a request (a CLI
+it after every other `teardown_appcontext` handler. A streamed response defers the close until the
+response is closed. Otherwise its lifetime is bounded by Flask's **app context**, not by Flask's
+request context: an app context pushed without a request (a CLI
 command, a bare `app.app_context()`) fires the same teardown with no child ever built, which is
 why the close is guarded.
 _Avoid_: request context — Flask's own object, which this is not; request scope — `Scope.REQUEST`

README.md

@@ -66,6 +66,8 @@ container.validate()  # optional fail-fast; must come after setup_di registers i
 
 Pass `auto_inject=True` to `setup_di` to wire every registered view (app and blueprint routes alike) without a per-view `@inject`. It walks `app.view_functions` when `setup_di` runs, so call `setup_di` after all routes are registered. `flask.Request` is resolvable within DI via the pre-built `flask_request_provider` context provider.
 
+A streamed response (a generator body, with or without `stream_with_context`) keeps the child container open until the response is closed, which a WSGI server does after sending the body. A `stream_with_context` generator can resolve from it, and dependencies resolved in the view stay usable while the body streams. A `direct_passthrough` response, such as one from `send_file`, is not closed through Flask, so its child container closes at teardown as usual.
+
 `async def` views work with `@inject` and with `auto_inject=True`. Flask runs them through asgiref, so install Flask's async extra (`pip install "flask[async]"`). The markers are resolved from the request's child container before the view is awaited. Teardown is sync: the child container is closed with `close_sync()`, so a `Scope.REQUEST` provider with an async finalizer is not supported and fails when the request ends. The same applies to the root container at shutdown: if an APP-scoped provider has an async finalizer, close the root with `asyncio.run(fetch_di_container(app).close_async())` instead of `close_sync()`.
 
 A container is open from construction. To fail fast at boot, call `.validate()` after `setup_di`, which registers `flask_request_provider`. Flask has no shutdown hook, so call `fetch_di_container(app).close_sync()` at process shutdown.

modern_di_flask/main.py

@@ -3,7 +3,7 @@
 import typing
 
 import flask
-from flask import Flask, Request, g
+from flask import Flask, Request, Response, g
 from modern_di import Container, Scope, integrations, providers
 
 
@@ -15,6 +15,7 @@
 _ROOT_CONTAINER_KEY = "modern_di_container"
 _CHILD_CONTAINER_ATTR = "modern_di_request_container"
 _CHILD_REQUEST_ATTR = "modern_di_request"
+_DEFERRED_CHILD_ATTR = "modern_di_deferred_request_container"
 _NO_CHILD_CONTAINER_MSG = (
     "No modern-di container found for this request. Call setup_di(app, container) "
     "so requests pass through the modern-di request hooks before using @inject."
@@ -53,9 +54,18 @@ def _enter_request() -> None:
     _request_child_container()
 
 
+def _defer_close_to_response(_sender: Flask, response: Response, **_extra: typing.Any) -> None:  # noqa: ANN401
+    child: Container | None = getattr(g, _CHILD_CONTAINER_ATTR, None)
+    if child is None or not response.is_streamed or response.direct_passthrough:
+        child = None
+    else:
+        response.call_on_close(child.close_sync)
+    setattr(g, _DEFERRED_CHILD_ATTR, child)
+
+
 def _close_request(_exception: BaseException | None) -> None:
     child: Container | None = getattr(g, _CHILD_CONTAINER_ATTR, None)
-    if child is not None:
+    if child is not None and child is not getattr(g, _DEFERRED_CHILD_ATTR, None):
         child.close_sync()
 
 
@@ -68,6 +78,7 @@ def setup_di(app: Flask, container: Container, *, auto_inject: bool = False) ->
     # Flask runs these in reverse, so putting the close first makes it run last.
     funcs = app.teardown_appcontext_funcs
     funcs[:] = [_close_request, *(func for func in funcs if func is not _close_request)]
+    flask.request_finished.connect(_defer_close_to_response, app)
     if auto_inject:
         _inject_views(app)
     return container

tests/test_request_hooks.py

@@ -114,28 +114,61 @@ def test_request_container_is_not_rebuilt_after_close(app: Flask) -> None:
     assert events == ["finalized"]
 
 
-def test_streamed_response_never_rebuilds_the_request_container(app: Flask) -> None:
+def test_streamed_response_keeps_the_request_container_until_the_response_closes(app: Flask) -> None:
     events: list[str] = []
-    outcomes: list[str] = []
 
     @app.route("/")
     def index() -> Response:
         session = _use()
 
         def generate() -> typing.Iterator[str]:
             yield "a"
-            try:
-                outcomes.append("same" if _use() is session else "rebuilt")
-            except exceptions.ContainerClosedError:
-                outcomes.append("closed")
+            events.append("same" if _use() is session else "rebuilt")
 
         return Response(stream_with_context(generate()))
 
     with Container(groups=[_make_group(events)]) as root:
         setup_di(app, root)
-        assert app.test_client().get("/").text == "a"
-    # Flask 3.0 keeps the context for the whole stream; 3.1.2+ tears it down before iterating.
-    assert outcomes in (["same"], ["closed"])
+        response = app.test_client().get("/")
+        assert response.text == "a"
+        assert events == ["same"]
+        response.close()
+    assert events == ["same", "finalized"]
+
+
+def test_plain_generator_response_keeps_resolved_dependencies_live(app: Flask) -> None:
+    events: list[str] = []
+
+    @app.route("/")
+    def index() -> Response:
+        _use()
+
+        def generate() -> typing.Iterator[str]:
+            yield "a"
+            events.append("streamed")
+
+        return Response(generate())
+
+    with Container(groups=[_make_group(events)]) as root:
+        setup_di(app, root)
+        response = app.test_client().get("/")
+        assert response.text == "a"
+        assert events == ["streamed"]
+        response.close()
+    assert events == ["streamed", "finalized"]
+
+
+def test_direct_passthrough_response_closes_at_teardown(app: Flask) -> None:
+    events: list[str] = []
+
+    @app.route("/")
+    def index() -> Response:
+        _use()
+        return Response(iter([b"a"]), direct_passthrough=True)
+
+    with Container(groups=[_make_group(events)]) as root:
+        setup_di(app, root)
+        app.test_client().get("/")
     assert events == ["finalized"]