Middleware for request header id - is this correct?

[bunny-therapist]

I tried to write middleware that, if the header "X-Request-Id" is in the request, then that same header with the same value will also be in the response. I think I got it right, but since the structure of these are quite complicated (and I sometimes see these lines in tracebacks - but I guess that could just be because all requests go thru it), I though I would paste it here and ask if I did everything correct.

I don't have the if-else for the scope that is in https://docs.litestar.dev/2/usage/middleware/creating-middleware.html but I assumed that was not necessary when setting the scope types on the middleware class to just the http scope.

I am on litestar 2.9.1

from typing import Optional

import litestar
import litestar.datastructures
import litestar.middleware
import litestar.types
from litestar.enums import ScopeType

X_REQUEST_ID_HEADER_NAME = "X-Request-Id"


class XRequestIdMiddleware(litestar.middleware.base.AbstractMiddleware):
    """If X-Request-Id is in the request headers, add it to the response headers.

    Reference: https://docs.litestar.dev/2/usage/middleware/creating-middleware.html

    """

    scopes = {ScopeType.HTTP}

    async def __call__(
        self,
        scope: litestar.types.Scope,
        receive: litestar.types.Receive,
        send: litestar.types.Send,
    ) -> None:
        async def send_wrapper(message: litestar.types.Message) -> None:
            if message["type"] == "http.response.start":
                request_id: Optional[str] = litestar.Request(scope).headers.get(
                    X_REQUEST_ID_HEADER_NAME
                )
                if request_id:
                    headers = litestar.datastructures.MutableScopeHeaders.from_message(
                        message=message
                    )
                    headers[X_REQUEST_ID_HEADER_NAME] = request_id
            await send(message)

        await self.app(scope, receive, send_wrapper)

[provinzkraut]

This looks good :)