fix(api): stop gzip middleware from dropping response starts on Starlette 1.8

invokeai/app/api_app.py

@@ -19,7 +19,7 @@
 from starlette.datastructures import Headers
 from starlette.exceptions import HTTPException as StarletteHTTPException
 from starlette.middleware.base import BaseHTTPMiddleware, RequestResponseEndpoint
-from starlette.middleware.gzip import GZipMiddleware, GZipResponder, IdentityResponder
+from starlette.middleware.gzip import GZipMiddleware
 from starlette.types import ASGIApp, Message, Receive, Scope, Send
 
 import invokeai.frontend.webv2 as webv2_dir
@@ -265,49 +265,51 @@ def _is_compressible(content_type: str) -> bool:
     return content_type.startswith(COMPRESSIBLE_CONTENT_TYPES)
 
 
-class _ContentTypeAwareGZipResponder(GZipResponder):
-    """Skips compression for response types that are already compressed.
-
-    `content_type_is_excluded` is computed when the response starts and only read once the
-    body arrives, so widening it right after the base class has set it is enough — no need to
-    reimplement Starlette's streaming/pathsend handling.
-    """
-
-    async def send_with_compression(self, message: Message) -> None:
-        await super().send_with_compression(message)
-        if message["type"] == "http.response.start" and not self.content_type_is_excluded:
-            self.content_type_is_excluded = not _is_compressible(
-                Headers(raw=message["headers"]).get("content-type", "")
-            )
-
-
-class ContentTypeAwareGZipMiddleware(GZipMiddleware):
+class ContentTypeAwareGZipMiddleware:
     """GZip, but only for content types that actually compress.
 
-    Starlette's GZipMiddleware compresses every response type except `text/event-stream`. The
-    gallery serves PNG, WebP and MP4 bytes, which are already compressed: a 3 MB PNG costs
-    ~52ms of event-loop time to gzip and comes back *larger* than it went in. With auto-switch
-    enabled the UI fetches the full image after every generated image, so that cost lands
-    repeatedly during a batch — exactly when the server can least afford to stall.
+    Older Starlette releases compress every response type except `text/event-stream`, and even
+    newer ones only skip a fixed list of media types. The gallery serves PNG, WebP and MP4 bytes,
+    which are already compressed: a 3 MB PNG costs ~52ms of event-loop time to gzip and comes
+    back *larger* than it went in. With auto-switch enabled the UI fetches the full image after
+    every generated image, so that cost lands repeatedly during a batch — exactly when the
+    server can least afford to stall.
+
+    Each response is routed when it starts: compressible ones through Starlette's GZipMiddleware,
+    everything else straight to the server. Only the middleware's public ASGI interface is used,
+    because FastAPI leaves Starlette unbounded above and the responder's internals change between
+    releases — Starlette 1.8 began forwarding `http.response.start` for its own excluded types
+    inside the responder, so patching its exclusion flag afterwards dropped the start message.
 
     Lowering `compresslevel` does not help here: on incompressible input, level 1 costs
     essentially the same as level 9 because deflate still has to scan the data. It does help a
     great deal on the compressible path — see `configure_gzip`.
     """
 
+    def __init__(self, app: ASGIApp, minimum_size: int, compresslevel: int) -> None:
+        self.app = app
+        self.minimum_size = minimum_size
+        self.compresslevel = compresslevel
+
     async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
         if scope["type"] != "http":
             await self.app(scope, receive, send)
             return
 
-        if "gzip" in Headers(scope=scope).get("Accept-Encoding", ""):
-            responder: ASGIApp = _ContentTypeAwareGZipResponder(
-                self.app, self.minimum_size, compresslevel=self.compresslevel
-            )
-        else:
-            responder = IdentityResponder(self.app, self.minimum_size)
+        async def route_by_content_type(scope: Scope, receive: Receive, gzip_send: Send) -> None:
+            target = send
+
+            async def routed_send(message: Message) -> None:
+                nonlocal target
+                if message["type"] == "http.response.start":
+                    content_type = Headers(raw=message.get("headers", [])).get("content-type", "")
+                    target = gzip_send if _is_compressible(content_type) else send
+                await target(message)
+
+            await self.app(scope, receive, routed_send)
 
-        await responder(scope, receive, send)
+        gzip = GZipMiddleware(route_by_content_type, minimum_size=self.minimum_size, compresslevel=self.compresslevel)
+        await gzip(scope, receive, send)
 
 
 # Responses below this are not worth a compression pass; the gzip framing alone is ~20 bytes.

tests/app/api/test_gzip_content_types.py

@@ -12,6 +12,7 @@
 
 import pytest
 from fastapi import FastAPI, Response
+from fastapi.responses import StreamingResponse
 from fastapi.testclient import TestClient
 
 from invokeai.app.api_app import ContentTypeAwareGZipMiddleware, configure_gzip
@@ -25,7 +26,12 @@ def _build_app(compresslevel: int) -> TestClient:
     app = FastAPI()
 
     @app.get("/payload")
-    def payload(content_type: str) -> Response:
+    def payload(content_type: str, stream: bool = False) -> Response:
+        if stream:
+            # Every response in the real app arrives like this: the BaseHTTPMiddleware
+            # subclasses re-emit bodies as `more_body` chunks.
+            chunks = [BODY[: len(BODY) // 2], BODY[len(BODY) // 2 :]]
+            return StreamingResponse(iter(chunks), media_type=content_type)
         return Response(content=BODY, media_type=content_type)
 
     @app.get("/tiny")
@@ -51,14 +57,18 @@ def client() -> TestClient:
         "image/svg+xml",
     ],
 )
-def test_compressible_types_are_compressed(client: TestClient, content_type: str):
-    r = client.get("/payload", params={"content_type": content_type}, headers={"Accept-Encoding": "gzip"})
+@pytest.mark.parametrize("stream", [False, True])
+def test_compressible_types_are_compressed(client: TestClient, content_type: str, stream: bool):
+    r = client.get(
+        "/payload", params={"content_type": content_type, "stream": stream}, headers={"Accept-Encoding": "gzip"}
+    )
 
     assert r.status_code == 200
     assert r.headers["content-encoding"] == "gzip"
     # httpx decodes transparently, so this also proves the compressed body round-trips.
     assert r.content == BODY
-    assert int(r.headers["content-length"]) < len(BODY)
+    if not stream:
+        assert int(r.headers["content-length"]) < len(BODY)
 
 
 @pytest.mark.parametrize(
@@ -69,14 +79,21 @@ def test_compressible_types_are_compressed(client: TestClient, content_type: str
         "image/jpeg",
         "video/mp4",
         "application/zip",
+        # Not excluded by Starlette itself, so only this middleware keeps them uncompressed. These
+        # are the types whose start message Starlette 1.8 lost under the previous implementation.
         "application/octet-stream",
+        "application/pdf",
+        "font/ttf",
         # `text/` prefix matches, but compressing an event stream withholds events until the
         # compressor flushes.
         "text/event-stream",
     ],
 )
-def test_already_compressed_types_are_passed_through(client: TestClient, content_type: str):
-    r = client.get("/payload", params={"content_type": content_type}, headers={"Accept-Encoding": "gzip"})
+@pytest.mark.parametrize("stream", [False, True])
+def test_already_compressed_types_are_passed_through(client: TestClient, content_type: str, stream: bool):
+    r = client.get(
+        "/payload", params={"content_type": content_type, "stream": stream}, headers={"Accept-Encoding": "gzip"}
+    )
 
     assert r.status_code == 200
     assert "content-encoding" not in r.headers, f"{content_type} should not be gzipped"
@@ -171,38 +188,6 @@ def test_out_of_range_levels_are_rejected(level: int):
         InvokeAIAppConfig(http_compression_level=level)
 
 
-def test_the_starlette_contract_the_responder_relies_on_still_holds():
-    """`_ContentTypeAwareGZipResponder` widens Starlette's own `content_type_is_excluded` right
-    after the base class has computed it. That is only safe while Starlette keeps exposing the
-    flag and keeps *buffering* `http.response.start` rather than forwarding it immediately. If an
-    upgrade changes either, the widening is silently lost and the app compresses PNGs again —
-    this test names the cause instead of leaving a puzzling content-type failure.
-    """
-    import asyncio
-
-    from starlette.middleware.gzip import IdentityResponder
-    from starlette.types import Message
-
-    responder = IdentityResponder(app=None, minimum_size=1)  # type: ignore[arg-type]
-    assert hasattr(responder, "content_type_is_excluded"), (
-        "Starlette no longer exposes `content_type_is_excluded`; the exclusion must be reimplemented"
-    )
-
-    forwarded: list[Message] = []
-
-    async def capture(message: Message) -> None:
-        forwarded.append(message)
-
-    responder.send = capture
-    start: Message = {"type": "http.response.start", "status": 200, "headers": [(b"content-type", b"image/png")]}
-    asyncio.run(responder.send_with_compression(start))
-
-    assert forwarded == [], (
-        "Starlette now forwards http.response.start immediately; the content type must be inspected "
-        "before delegating to the base responder"
-    )
-
-
 def test_clients_without_gzip_support_get_plain_bodies(client: TestClient):
     r = client.get("/payload", params={"content_type": "application/json"}, headers={"Accept-Encoding": "identity"})