From 07ec64c66553dd34656bdbad041346a4cf4c0bd3 Mon Sep 17 00:00:00 2001 From: CyberneticX-Tech <138270143+CyberneticX-Tech@users.noreply.github.com> Date: Mon, 7 Sep 2026 14:47:12 +0000 Subject: [PATCH] fix(litestar): register plugin middleware at the outer boundary The plugin appended its middleware to app_config.middleware, which makes them the innermost entries in the ASGI pipeline. A request rejected by user middleware that runs earlier (auth, CORS, rate limiting) never reached the correlation middleware, so its access log and error hooks had no correlation ID. Prepend instead so the correlation context is established before user middleware runs. Closes #729 --- sqlspec/extensions/litestar/plugin.py | 6 +++- .../litestar/test_correlation_middleware.py | 32 ++++++++++++++++++- 2 files changed, 36 insertions(+), 2 deletions(-) diff --git a/sqlspec/extensions/litestar/plugin.py b/sqlspec/extensions/litestar/plugin.py index 8934872cf..9061e3720 100644 --- a/sqlspec/extensions/litestar/plugin.py +++ b/sqlspec/extensions/litestar/plugin.py @@ -399,7 +399,11 @@ def store_sqlspec_in_state() -> None: if self._enable_sqlcommenter_middleware: new_middlewares.append(DefineMiddleware(SQLCommenterMiddleware)) if new_middlewares: - app_config.middleware = [*(app_config.middleware or []), *new_middlewares] + # Prepend so correlation context is established at the outer boundary of the + # pipeline. Appending makes these the innermost middleware, so a request + # rejected earlier (auth, CORS, rate limiting) never reaches them and its + # access log and error hooks have no correlation ID. + app_config.middleware = [*new_middlewares, *(app_config.middleware or [])] log_with_context( logger, diff --git a/tests/integration/extensions/litestar/test_correlation_middleware.py b/tests/integration/extensions/litestar/test_correlation_middleware.py index 380256cc9..c1438c371 100644 --- a/tests/integration/extensions/litestar/test_correlation_middleware.py +++ b/tests/integration/extensions/litestar/test_correlation_middleware.py @@ -1,6 +1,8 @@ -from typing import Any, cast +from typing import TYPE_CHECKING, Any, cast from litestar import Litestar, get +from litestar.middleware import AbstractMiddleware, DefineMiddleware +from litestar.response.base import ASGIResponse from litestar.testing import TestClient from sqlspec import SQLSpec @@ -9,6 +11,9 @@ from sqlspec.extensions.litestar import SQLSpecPlugin from sqlspec.utils.correlation import CorrelationContext +if TYPE_CHECKING: + from litestar.types import Receive, Scope, Send + def setup_function() -> None: """Clear correlation context before each test to prevent pollution.""" @@ -99,3 +104,28 @@ def test_correlation_middleware_auto_detection_can_be_disabled() -> None: response = client.get("/correlation", headers={"X-Custom-ID": "custom-value"}) assert response.json()["correlation_id"] == "custom-value" + + +def test_correlation_middleware_runs_before_user_middleware() -> None: + """A request rejected by user middleware still carries correlation context.""" + observed: dict[str, str | None] = {} + + class RejectingMiddleware(AbstractMiddleware): + async def __call__(self, scope: "Scope", receive: "Receive", send: "Send") -> None: + observed["correlation_id"] = CorrelationContext.get() + await ASGIResponse(body=b'{"detail":"unauthorized"}', status_code=401)(scope, receive, send) + + extension_config = cast("ExtensionConfigs", {"litestar": {"enable_correlation_middleware": True}}) + spec = SQLSpec() + spec.add_config(SqliteConfig(connection_config={"database": ":memory:"}, extension_config=extension_config)) + app = Litestar( + route_handlers=[correlation_handler], + plugins=[SQLSpecPlugin(sqlspec=spec)], + middleware=[DefineMiddleware(RejectingMiddleware)], + ) + + with TestClient(app) as client: + response = client.get("/correlation", headers={"X-Request-ID": "abc-123"}) + + assert response.status_code == 401 + assert observed["correlation_id"] == "abc-123"