life-echo/api/app/core/middleware.py

"""
HTTP 中间件：request_id 注入。
"""

import uuid

from starlette.datastructures import State
from starlette.requests import Request
from starlette.types import ASGIApp, Message, Receive, Scope, Send

from app.core.logging import logger
from app.core.telemetry import current_trace_context


class RequestIdMiddleware:
    """Inject request_id into request.state and response headers, bind to loguru context.

    Pure ASGI middleware (not BaseHTTPMiddleware) so FastAPI exception handlers still run.
    """

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

    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
        if scope["type"] != "http":
            await self.app(scope, receive, send)
            return

        request_id = None
        for name, value in scope.get("headers", []):
            if name == b"x-request-id":
                request_id = value.decode("latin-1")
                break
        if not request_id:
            request_id = str(uuid.uuid4())
        existing_state = scope.get("state")
        if isinstance(existing_state, State):
            existing_state.request_id = request_id
        elif isinstance(existing_state, dict):
            existing_state["request_id"] = request_id
            scope["state"] = State(existing_state)
        else:
            scope["state"] = State({"request_id": request_id})

        bind = {"request_id": request_id, **current_trace_context()}

        async def send_with_request_id(message: Message) -> None:
            if message["type"] == "http.response.start":
                headers = list(message.get("headers", []))
                headers.append((b"x-request-id", request_id.encode("latin-1")))
                message = {**message, "headers": headers}
            await send(message)

        with logger.contextualize(**bind):
            await self.app(scope, receive, send_with_request_id)