brain-of-reese/app/core/compression.py

"""Selective gzip for every HTTP response but the SSE chat stream (phase 167, D1, TODO L5).

The defect (TODO L5): ``GET /api/docs/tree`` loads ~2 MB in production and
the Sources page pays that on every sync and every shell boot. Nothing in
this stack compressed — ``app/main.py`` added only Session / caching /
security-headers, and no response ever carried ``Content-Encoding``. KB
JSON is extremely repetitive (key names, paths, ISO timestamps), so
gzipping it lands at roughly **8-12× smaller** on the wire: the same bytes
the client already knows how to read, ~one order of magnitude faster.

The SSE stream is EXCLUDED (A15, phase 167 LOCKED A1): ``POST /api/chat``
is a ``text/event-stream`` whose whole value is that frames arrive *while
the model is still talking*. Any middleware that buffers the response —
``BaseHTTPMiddleware`` is the classic offender, and the reason this module
is **pure ASGI** — would deliver the entire turn at the end, silently
breaking the typing animation, the tool lines, and the cancel surface. So
the skip is by EXACT path (never a prefix: ``/api/chat-images/…`` keeps
its compression) and the skipped scope is handed to the app untouched.

The compression math is deliberately NOT re-implemented here. Starlette's
:class:`~starlette.middleware.gzip.GZipMiddleware` owns ``Accept-Encoding``
negotiation (a client that does not speak gzip gets identity bytes), the
``minimum_size`` floor, the skip when a response ALREADY sets
``content-encoding``, the ``Vary: Accept-Encoding`` append, and the
streaming ``compressobj`` + final flush. It is also what makes the stack
cooperate with a fronting proxy: **nginx may also gzip, but an
already-encoded body is never double-encoded** — Starlette preserves the
upstream ``Content-Encoding`` verbatim, and so does nginx.

One seam IS patched here, and it is framing, not math: the
``CachingMiddleware`` below us is a ``BaseHTTPMiddleware``, and starlette
1.x re-frames every response it emits upstream (``_StreamingResponse``
ships each body chunk with ``more_body=True`` plus an empty terminator),
so Starlette's buffered small-response exemption can never fire for
anything that travelled through it — a 33-byte ``/api/health`` would
stream into the compressor and come back INFLATED, exactly what the
``minimum_size`` floor exists to prevent. :class:`_RebufferSmallResponses`
sits under the responder and hands it back the single buffered body
message the exemption was written for, bounded to the floor.
"""
from __future__ import annotations

from starlette.datastructures import Headers
from starlette.middleware.gzip import GZipMiddleware
from starlette.types import ASGIApp, Message, Receive, Scope, Send

#: Paths that are NEVER compressed, matched EXACTLY (never by prefix).
#: ``POST /api/chat`` is the SSE stream (A15): frame-by-frame delivery is
#: the entire point of the endpoint, and buffering the response to feed it
#: through a compressor would deliver the whole turn at completion (and a
#: never-ending stream never completes on a cancel). The exact match is
#: load-bearing — ``/api/chat-images/…`` shares the prefix and must keep
#: routing through the compression layer.
DEFAULT_SKIP_PATHS: tuple[str, ...] = ("/api/chat",)


class _RebufferSmallResponses:
    """Hold a declared-small body so Starlette's floor exemption fires.

    When ``http.response.start`` carries ``content-length`` below
    ``minimum_size`` — byte-for-byte the condition Starlette's buffered
    path checks — the messages are held (bounded by the floor) and
    re-emitted as ONE ``http.response.body`` with ``more_body=False``,
    undoing the ``BaseHTTPMiddleware`` re-framing so the responder's
    small-response exemption sees what it was written to see and the
    bytes stay identity (no ``Content-Encoding``, no ``Vary`` append).

    A lying header can only cost the buffer, never correctness: if the
    body outgrows the floor despite the header, everything held is
    forwarded as it came and the rest streams through the normal
    (compressing) path. Non-body messages while holding (e.g.
    ``http.response.pathsend``) also flush the hold untouched.
    """

    def __init__(self, app: ASGIApp, *, minimum_size: int) -> None:
        self.app = app
        self.minimum_size = minimum_size

    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
        holding = False
        armed = True  # the start decision happens at most once per response
        held: list[Message] = []
        total = 0

        async def send_filter(message: Message) -> None:
            nonlocal holding, armed, total
            message_type = message["type"]

            if message_type == "http.response.start" and armed:
                armed = False
                length = Headers(scope=message).get("content-length")
                if (
                    length is not None
                    and length.isdigit()
                    and int(length) < self.minimum_size
                ):
                    held.append(message)
                    holding = True
                    return
                await send(message)
                return

            if holding and message_type == "http.response.body":
                held.append(message)
                total += len(message.get("body", b""))
                if not message.get("more_body", False):
                    # Whole body in hand, still under the floor: re-emit
                    # it as the single buffered message the exemption wants.
                    start = held[0]
                    merged = b"".join(m.get("body", b"") for m in held[1:])
                    holding = False
                    held.clear()
                    await send(start)
                    await send({"type": "http.response.body", "body": merged, "more_body": False})
                elif total >= self.minimum_size:
                    # The content-length lied — stop holding, stream on.
                    holding = False
                    for held_message in held:
                        await send(held_message)
                    held.clear()
                return

            if holding:
                # pathsend or anything unexpected mid-hold: flush as-is.
                holding = False
                for held_message in held:
                    await send(held_message)
                held.clear()

            await send(message)

        await self.app(scope, receive, send_filter)


class SelectiveGZipMiddleware:
    """Pure-ASGI selective gzip (phase 167, D1 — no ``BaseHTTPMiddleware``).

    Holds the wrapped ``app`` and a composed
    :class:`~starlette.middleware.gzip.GZipMiddleware` built over the SAME
    app (via the :class:`_RebufferSmallResponses` seam, so the floor still
    works under the caching layer's re-framing); ``__call__`` delegates to
    the gzip path unless the scope is not HTTP (websocket / lifespan pass
    straight through) or the request path is one of ``skip_paths`` — in
    which case the app is called DIRECTLY, with the original
    ``receive``/``send``, so a streaming response reaches the client frame
    by frame, uncompressed.

    ``minimum_size`` is the Starlette floor (bytes): responses smaller than
    it are forwarded without compression because gzip can make them
    bigger. It is supplied by the caller (``Settings.compress_min_bytes``;
    ``0`` removes this middleware entirely at wiring time — see
    ``app/main.py``). The SSE stream (``POST /api/chat``, A15) is never
    compressed at any value.
    """

    def __init__(
        self,
        app: ASGIApp,
        *,
        minimum_size: int,
        skip_paths: tuple[str, ...] = DEFAULT_SKIP_PATHS,
    ) -> None:
        self.app = app
        self.minimum_size = minimum_size
        self.skip_paths = tuple(skip_paths)
        # The compression itself stays Starlette's: negotiation, floor,
        # already-encoded skip, Vary append, streaming compressobj + flush.
        self.gzip = GZipMiddleware(
            _RebufferSmallResponses(app, minimum_size=minimum_size),
            minimum_size=minimum_size,
        )

    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
        if scope["type"] != "http" or scope["path"] in self.skip_paths:
            # Non-HTTP scopes have no response to compress; the SSE stream
            # must never see this layer (A15 — see DEFAULT_SKIP_PATHS).
            await self.app(scope, receive, send)
            return
        await self.gzip(scope, receive, send)