feat(middleware): add W3C Trace Context correlation middleware

docs/examples/middleware/correlation_standalone.py

@@ -0,0 +1,13 @@
+from litestar import Litestar, Request, get
+from litestar.middleware.correlation import CorrelationMiddleware, get_correlation_id
+
+
+@get("/")
+async def index_handler(request: Request) -> dict[str, str | None]:
+    return {"correlation_id": get_correlation_id(request)}
+
+
+app = Litestar(
+    route_handlers=[index_handler],
+    middleware=[CorrelationMiddleware()],
+)

docs/reference/middleware/correlation.rst

@@ -0,0 +1,6 @@
+===========
+correlation
+===========
+
+.. automodule:: litestar.middleware.correlation
+    :members:

docs/reference/middleware/index.rst

@@ -10,6 +10,7 @@ middleware
     allowed_hosts
     authentication
     compression
+    correlation
     csrf
     logging
     rate_limit

docs/release-notes/changelog.rst

@@ -13,6 +13,16 @@
         dynamic ``LISTEN`` and ``UNLISTEN`` operations no longer contend with the listener.
         Psycopg 3.2.4 or newer is now required for development and documentation builds.
 
+    .. change:: Add W3C Trace Context Correlation Middleware
+        :type: feature
+        :issue: 4719
+
+        Added standalone W3C Trace Context correlation middleware via
+        :class:`~litestar.middleware.correlation.CorrelationMiddleware`.
+
+        .. seealso::
+            :doc:`/usage/middleware/correlation`
+
     .. change:: Move ``httpx`` to the ``testing`` extra
         :type: feature
         :pr: 4950

docs/usage/middleware/correlation.rst

@@ -0,0 +1,72 @@
+======================
+Correlation Middleware
+======================
+
+The correlation middleware extracts, generates, and propagates correlation and trace IDs for each connection.
+
+This facilitates distributed tracing and unified logging across microservices and async handlers.
+
+Features
+--------
+
+- Priority header fallback (:code:`x-request-id`, :code:`x-correlation-id`, and :code:`traceparent` by default).
+- W3C :code:`traceparent` defensive parsing.
+- UUID4 generation fallback when no header matches.
+- The active ID is stored on the connection scope, making it available to handlers and other middlewares.
+- Optional response-header propagation.
+
+Usage
+-----
+
+.. literalinclude:: /examples/middleware/correlation_standalone.py
+    :language: python
+
+Accessing the correlation ID
+----------------------------
+
+The active correlation ID is stored on the connection scope and can be retrieved anywhere the scope is available -
+in handlers, dependencies, or other middlewares - using
+:func:`~litestar.middleware.correlation.get_correlation_id`:
+
+.. code-block:: python
+
+    from litestar import Request, get
+    from litestar.middleware.correlation import get_correlation_id
+
+
+    @get("/")
+    async def handler(request: Request) -> str | None:
+        return get_correlation_id(request)
+
+The helper also accepts a raw ASGI scope or a :class:`~litestar.connection.WebSocket` directly:
+
+.. code-block:: python
+
+    from litestar import WebSocket, websocket
+
+
+    @websocket("/ws")
+    async def websocket_handler(socket: WebSocket) -> None:
+        await socket.accept()
+        await socket.send_text(get_correlation_id(socket) or "missing")
+
+Header behavior
+---------------
+
+The middleware validates W3C ``traceparent`` values. Values from all other configured headers are treated as opaque
+correlation values. Additional formats, including ``grpc-trace-bin`` and provider-specific headers, can be selected
+without repeating the defaults:
+
+.. code-block:: python
+
+    CorrelationMiddleware(
+        additional_header_names=("grpc-trace-bin", "x-cloud-trace-context"),
+    )
+
+Use ``header_names`` instead to replace the complete lookup list and control its priority. The two options are
+mutually exclusive. The middleware does not parse additional or replacement header formats.
+
+Incoming scope headers are not modified. By default, the selected correlation ID replaces ``x-request-id`` in the
+response; set ``response_header_name=None`` to disable this. This response header contains the selected correlation
+value, not a raw copy of the incoming header. The middleware does not propagate correlation headers to outbound
+requests.

docs/usage/middleware/index.rst

@@ -19,4 +19,5 @@ See :doc:`the documentation regarding these </usage/middleware/builtin-middlewar
 
     using-middleware
     builtin-middleware
+    correlation
     creating-middleware

litestar/middleware/__init__.py

@@ -8,12 +8,20 @@
     DefineMiddleware,
     MiddlewareProtocol,
 )
+from litestar.middleware.correlation import (
+    TRACE_CONTEXT_FALLBACK_HEADERS,
+    CorrelationMiddleware,
+    get_correlation_id,
+)
 
 __all__ = (
+    "TRACE_CONTEXT_FALLBACK_HEADERS",
     "ASGIMiddleware",
     "AbstractAuthenticationMiddleware",
     "AbstractMiddleware",
     "AuthenticationResult",
+    "CorrelationMiddleware",
     "DefineMiddleware",
     "MiddlewareProtocol",
+    "get_correlation_id",
 )

litestar/middleware/correlation.py

@@ -0,0 +1,184 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any, Final
+from uuid import uuid4
+
+from litestar.connection import ASGIConnection
+from litestar.datastructures.headers import Headers, MutableScopeHeaders
+from litestar.enums import ScopeType
+from litestar.middleware.base import ASGIMiddleware
+from litestar.types import Empty
+from litestar.utils.scope.state import ScopeState
+
+if TYPE_CHECKING:
+    from collections.abc import Sequence
+
+    from litestar.types import ASGIApp, Message, Receive, Scope, Send
+
+__all__ = (
+    "TRACE_CONTEXT_FALLBACK_HEADERS",
+    "CorrelationMiddleware",
+    "get_correlation_id",
+)
+
+TRACE_CONTEXT_FALLBACK_HEADERS: Final[tuple[str, ...]] = (
+    "x-request-id",
+    "x-correlation-id",
+    "traceparent",
+)
+
+_LOWERCASE_HEX = frozenset("0123456789abcdef")
+
+
+def get_correlation_id(connection: ASGIConnection[Any, Any, Any, Any] | Scope) -> str | None:
+    """Get the correlation ID stored on the connection scope by :class:`CorrelationMiddleware`.
+
+    Args:
+        connection: An ASGI connection or scope.
+
+    Returns:
+        The correlation ID, or ``None`` if none was set.
+    """
+    scope = connection.scope if isinstance(connection, ASGIConnection) else connection
+    correlation_id = ScopeState.from_scope(scope).correlation_id
+    return None if correlation_id is Empty else correlation_id
+
+
+class CorrelationMiddleware(ASGIMiddleware):
+    """ASGI middleware for extracting, generating, and propagating correlation IDs.
+
+    The active correlation ID is stored on the connection scope and can be retrieved
+    with :func:`get_correlation_id`.
+    """
+
+    scopes = (ScopeType.HTTP, ScopeType.WEBSOCKET)
+
+    def __init__(
+        self,
+        header_names: Sequence[str] | None = None,
+        additional_header_names: Sequence[str] | None = None,
+        response_header_name: str | None = "x-request-id",
+        max_length: int = 128,
+    ) -> None:
+        """Initialize CorrelationMiddleware.
+
+        Args:
+            header_names: Header name or sequence of header names to inspect in priority order, replacing the defaults.
+            additional_header_names: Header name or sequence of header names to inspect after the defaults.
+            response_header_name: Optional header name to echo correlation ID in response. Set to None to disable.
+            max_length: Maximum length for correlation IDs to prevent log injection.
+
+        Raises:
+            ValueError: If ``max_length`` is not positive or both header name options are provided.
+        """
+        if max_length <= 0:
+            raise ValueError("max_length must be greater than 0")
+        if header_names is not None and additional_header_names is not None:
+            raise ValueError("header_names and additional_header_names are mutually exclusive")
+        if header_names is None:
+            if isinstance(additional_header_names, str):
+                additional_header_names = (additional_header_names,)
+            header_names = (*TRACE_CONTEXT_FALLBACK_HEADERS, *(additional_header_names or ()))
+        if isinstance(header_names, str):
+            header_names = (header_names,)
+        normalized_header_names: list[str] = []
+        for name in header_names:
+            normalized_name = name.strip().casefold()
+            if normalized_name and normalized_name not in normalized_header_names:
+                normalized_header_names.append(normalized_name)
+        self.header_names = tuple(normalized_header_names)
+        self.response_header_name = response_header_name.strip().casefold() if response_header_name else None
+        self.max_length = max_length
+
+    async def handle(self, scope: Scope, receive: Receive, send: Send, next_app: ASGIApp) -> None:
+        """ASGI call handler.
+
+        Args:
+            scope: The ASGI scope.
+            receive: The ASGI receive callable.
+            send: The ASGI send callable.
+            next_app: The next ASGI application in the middleware stack.
+        """
+        correlation_id = self._extract_correlation_id(scope)
+        ScopeState.from_scope(scope).correlation_id = correlation_id
+
+        if (response_header_name := self.response_header_name) is None:
+            await next_app(scope, receive, send)
+            return
+
+        async def send_wrapper(message: Message) -> None:
+            if message["type"] == "http.response.start":
+                headers = MutableScopeHeaders.from_message(message)
+                headers[response_header_name] = correlation_id
+            await send(message)
+
+        await next_app(scope, receive, send_wrapper)
+
+    def _extract_correlation_id(self, scope: Scope) -> str:
+        """Extract correlation ID from incoming request headers or generate fallback.
+
+        Args:
+            scope: The ASGI scope.
+
+        Returns:
+            Extracted or generated correlation ID.
+        """
+        headers = Headers.from_scope(scope)
+        for name in self.header_names:
+            if (value := headers.get(name)) is None:
+                continue
+            correlation_id = self._parse_traceparent(value) if name == "traceparent" else self._sanitize(value)
+            if correlation_id is not None:
+                return correlation_id
+        return str(uuid4())
+
+    def _parse_traceparent(self, value: str) -> str | None:
+        """Defensively parse W3C traceparent header.
+
+        Args:
+            value: Incoming traceparent header value.
+
+        Returns:
+            The extracted trace ID if valid, else the sanitized raw header string.
+        """
+        sanitized = _strip_safe_value(value)
+        if sanitized is None:
+            return None
+        parts = sanitized.split("-")
+        if len(parts) != 4:
+            return sanitized[: self.max_length]
+        version, trace_id, parent_id, flags = parts
+        if (
+            _is_lowercase_hex(version, 2)
+            and version != "ff"
+            and _is_lowercase_hex(trace_id, 32)
+            and trace_id != "0" * 32
+            and _is_lowercase_hex(parent_id, 16)
+            and parent_id != "0" * 16
+            and _is_lowercase_hex(flags, 2)
+        ):
+            return trace_id[: self.max_length]
+        return sanitized[: self.max_length]
+
+    def _sanitize(self, value: str) -> str | None:
+        """Sanitize a correlation ID by stripping whitespace and rejecting control characters.
+
+        Args:
+            value: Raw correlation ID.
+
+        Returns:
+            Sanitized correlation ID, or ``None`` when the value is unsafe.
+        """
+        sanitized = _strip_safe_value(value)
+        return sanitized[: self.max_length] if sanitized is not None else None
+
+
+def _is_lowercase_hex(value: str, length: int) -> bool:
+    return len(value) == length and _LOWERCASE_HEX.issuperset(value)
+
+
+def _strip_safe_value(value: str) -> str | None:
+    value = value.strip()
+    if not value or any(ord(character) < 32 or ord(character) == 127 for character in value):
+        return None
+    return value

litestar/utils/scope/state.py

@@ -31,6 +31,7 @@ class ScopeState:
         "body",
         "content_type",
         "cookies",
+        "correlation_id",
         "csrf_token",
         "dependency_cache",
         "do_cache",
@@ -55,6 +56,7 @@ def __init__(self) -> None:
         self.body = Empty
         self.content_type = Empty
         self.cookies = Empty
+        self.correlation_id = Empty
         self.csrf_token = Empty
         self.dependency_cache = Empty
         self.do_cache = Empty
@@ -77,6 +79,7 @@ def __init__(self) -> None:
     body: bytes | EmptyType
     content_type: tuple[str, dict[str, str]] | EmptyType
     cookies: dict[str, str] | EmptyType
+    correlation_id: str | EmptyType
     csrf_token: str | EmptyType
     dependency_cache: dict[str, Any] | EmptyType
     do_cache: bool | EmptyType