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

"""The byte budget: ONE log line per ``/api/*`` response (phase 167, task 05, D5, LOCKED A5).

TODO L5 asked "which endpoint is heavy, right now, in production?" — and
before this phase the app could not answer it. The 2 MB
``GET /api/docs/tree`` defect was found by a human eyeballing a browser
network tab; every other candidate endpoint was a guess. This module makes
the answer a ``grep`` away:

.. code-block:: text

    api method=GET path=/api/docs/tree status=200 bytes=204511 enc=gzip duration_ms=41

* ``bytes`` is the size of the body the ROUTE produced — the sum of the
  ``http.response.body`` chunk lengths as they pass through, uncompressed
  even when the wire carried gzip (the compression layer above us encodes
  AFTER we counted). ``enc`` reports what actually went on the wire
  (``content-encoding``, else ``identity``). The two together answer "how
  fat is it, and how fat did the client see it".
* level: ``INFO`` when ``bytes >= settings.log_large_response_bytes``
  (``BOR_LOG_LARGE_RESPONSE_BYTES``, default 65536; ``0`` = every ``/api/*``
  response at INFO), ``DEBUG`` below it — so the default log stays quiet
  and an offender is loud without a re-deploy.
* a ``304`` logs ``bytes=0`` deliberately: the revalidation win this phase
  bought is visible in the log, not hidden behind a missing line.
* `` stream=1`` is appended for ``text/event-stream`` responses, where
  ``bytes`` is the streamed total and ``duration_ms`` is the stream's LIFE.

Why **pure ASGI** (the phase-167 rule for anything on this path): the
middleware must never buffer. It only reads message lengths, so the SSE
chat stream (A15) passes through byte-identical — a
``BaseHTTPMiddleware`` here would deliver the whole turn at the end,
which is exactly the failure mode phase 167 exists to avoid elsewhere.

Logged on the ``http.response.body`` message with ``more_body=False`` —
for a stream that is when the **stream ends**, so ``duration_ms`` is the
stream's whole life. A stream the client walks away from (a cancelled SSE
turn) writes exactly one line too, with the bytes that actually got out:
the app unwinds without a completing body message and the ``finally`` path
emits it. Half a response happened, so half a count is the honest number —
and no line at all would make a chatty, abandoned poll look free.

An app exception is handled the same way (one line, bytes as counted, status
500 when the app never got to answer) and re-raises UNCHANGED; non-HTTP
scopes and non-``/api/`` paths pass through silently (pages, assets and the
static catch-all are not the catalogue surface this budget governs).
"""
from __future__ import annotations

import logging
import time

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

from app.config import get_settings

#: The logger behind every line — renders as ``… :: api method=…`` with the
#: §9 format, so ``grep ' :: api '`` isolates the budget log.
logger = logging.getLogger("app.response_budget")

#: Only paths under this prefix are measured. The admin catalogue (and
#: whatever new read route joins it) is the surface TODO L5's "any any
#: other endpoints that could load a ton of data" is about; pages, assets
#: and the static catch-all are not.
API_PATH_PREFIX = "/api/"

#: The media type whose line carries `` stream=1`` (A15 SSE): the frame
#: stream is where ``bytes`` is a total over a life rather than one body.
STREAM_MEDIA_TYPE = "text/event-stream"

#: Status logged when the app died before it could send
#: ``http.response.start`` — the exception is propagating and Starlette's
#: ``ServerErrorMiddleware`` (outermost) answers exactly that.
UNKNOWN_STATUS = 500


class _Tally:
    """What one measured response accumulates: status, the start message,
    bytes.

    ``start`` holds the ``http.response.start`` MESSAGE (the dict), not a
    snapshot of its headers, and its headers are read at emit time — after
    the final body chunk has been awaited. That is how ``enc`` reports the
    wire encoding from a layer sitting UNDER the compressor (the position
    that gives an uncompressed ``bytes=``): Starlette's gzip responder holds
    that same start message while it decides what to do with the body and
    decorates it in place (``MutableHeaders(raw=initial_message["headers"])``
    shares the list), so by the time the body completes the
    ``content-encoding`` the client sees is on it. Holding the message
    rather than the list is not pedantry — ``Headers(scope=message)``
    REBINDS ``message["headers"]`` to a fresh list (the compression
    responder and this stack's ``_RebufferSmallResponses`` both call it
    that way), so a captured list would be a stale one.

    The dependency is documented and pinned, not assumed:
    ``tests/integration/test_response_budget_api.py`` reads the real
    middleware stack for both ``enc=gzip`` and ``enc=identity``, so a
    future layer that copies the message instead of decorating it fails a
    test rather than quietly mislabeling the log (the byte count stays
    exact either way — it is the number the budget compares).
    """

    __slots__ = ("bytes", "logged", "start", "status")

    def __init__(self) -> None:
        self.bytes = 0
        self.status = UNKNOWN_STATUS
        self.start: Message | None = None
        self.logged = False

    def line(self, scope: Scope, started: float) -> str:
        """The §9-style single line — the exact shape README documents.

        ``duration_ms`` is rounded from the monotonic clock (never
        wall-clock: an NTP step must not produce a negative duration).
        """
        headers = Headers(raw=list(self.start["headers"]) if self.start else [])
        enc = (headers.get("content-encoding") or "identity").strip() or "identity"
        media = (headers.get("content-type") or "").partition(";")[0].strip().lower()
        duration_ms = int(round((time.monotonic() - started) * 1000))
        line = (
            f"api method={scope.get('method', '')} path={scope.get('path', '')} "
            f"status={self.status} bytes={self.bytes} enc={enc} "
            f"duration_ms={duration_ms}"
        )
        if media == STREAM_MEDIA_TYPE:
            line += " stream=1"
        return line

    def emit(self, scope: Scope, started: float, budget: int) -> None:
        """Write the line ONCE, at the level the budget decides.

        Idempotent on purpose: the completion path and the ``finally``
        path describe the same response and must never double-log it.
        """
        if self.logged:
            return
        self.logged = True
        line = self.line(scope, started)
        if self.bytes >= budget:
            logger.info(line)
        else:
            logger.debug(line)


class ResponseBudgetMiddleware:
    """Pure-ASGI response-size logger — zero buffering, SSE-safe.

    Wraps ``send`` only (never ``receive``): the timing runs from
    ``time.monotonic()`` at entry to the ``more_body=False`` completion,
    so ``duration_ms`` is the response's (or the stream's) whole life
    without the middleware needing to see one inbound message. Every
    ``http.response.body`` chunk's length adds to the count and nothing is
    held, so a stream still flows frame by frame.

    ``threshold`` overrides the INFO/DEBUG cut for this instance; the
    default (``None``) reads ``Settings.log_large_response_bytes`` per
    response, so ``app/main.py`` wires it with no arguments and
    ``BOR_LOG_LARGE_RESPONSE_BYTES`` stays the operator's knob (``0`` = log
    every ``/api/*`` response at INFO, negative refuses startup in the
    validator). LOCKED A5: this is a log line and a script, never a 413 —
    no endpoint starts rejecting large payloads here.
    """

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

    def _budget(self) -> int:
        if self.threshold is not None:
            return self.threshold
        return get_settings().log_large_response_bytes

    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
        if scope["type"] != "http" or not str(scope.get("path", "")).startswith(
            API_PATH_PREFIX
        ):
            # Nothing to measure: not HTTP (websocket/lifespan) or not the
            # API surface. Silent pass-through — the original send.
            await self.app(scope, receive, send)
            return

        started = time.monotonic()
        tally = _Tally()

        async def send_counting(message: Message) -> None:
            message_type = message["type"]
            if message_type == "http.response.start":
                tally.status = int(message["status"])
                # The MESSAGE, not a header snapshot — see _Tally: the
                # layers above decorate this same dict in place, so reading
                # it after the body completes yields the wire encoding.
                tally.start = message
            elif message_type == "http.response.body":
                tally.bytes += len(message.get("body", b""))
                if not message.get("more_body", False):
                    # Await FIRST: the layers above decorate the start
                    # headers while handling this final chunk, so `enc`
                    # reads what actually went on the wire.
                    await send(message)
                    tally.emit(scope, started, self._budget())
                    return
            await send(message)

        try:
            await self.app(scope, receive, send_counting)
        finally:
            # A normal response already emitted at its completion; this
            # covers an app that raised (or returned) without finishing a
            # body. The exception itself re-raises UNCHANGED.
            tally.emit(scope, started, self._budget())