diff --git a/litestar/middleware/logging.py b/litestar/middleware/logging.py index c986eb6433..1f73c7079c 100644 --- a/litestar/middleware/logging.py +++ b/litestar/middleware/logging.py @@ -191,7 +191,9 @@ def extract_response_data(self, scope: Scope) -> dict[str, Any]: connection_state = ScopeState.from_scope(scope) extracted_data = self.response_extractor( messages=( - connection_state.log_context.pop(HTTP_RESPONSE_START), + # NOTE: we don't pop the start message from the logging context in case + # there are multiple body messages to be logged + connection_state.log_context[HTTP_RESPONSE_START], connection_state.log_context.pop(HTTP_RESPONSE_BODY), ), ) @@ -224,6 +226,10 @@ async def send_wrapper(message: Message) -> None: elif message["type"] == HTTP_RESPONSE_BODY: connection_state.log_context[HTTP_RESPONSE_BODY] = message self.log_response(scope=scope) + + if not message["more_body"]: + connection_state.log_context.clear() + await send(message) return send_wrapper diff --git a/tests/e2e/test_middleware/test_logging_middleware_with_multi_body_response.py b/tests/e2e/test_middleware/test_logging_middleware_with_multi_body_response.py new file mode 100644 index 0000000000..0b9c2ef2e3 --- /dev/null +++ b/tests/e2e/test_middleware/test_logging_middleware_with_multi_body_response.py @@ -0,0 +1,30 @@ +from litestar import asgi +from litestar.middleware.logging import LoggingMiddlewareConfig +from litestar.testing import create_async_test_client +from litestar.types.asgi_types import Receive, Scope, Send + + +@asgi("/") +async def asgi_app(scope: Scope, receive: Receive, send: Send) -> None: + await send( + { + "type": "http.response.start", + "status": 200, + "headers": [ + (b"content-type", b"text/event-stream"), + (b"cache-control", b"no-cache"), + (b"connection", b"keep-alive"), + ], + } + ) + + # send two bodies + await send({"type": "http.response.body", "body": b"data: 1\n", "more_body": True}) + await send({"type": "http.response.body", "body": b"data: 2\n", "more_body": False}) + + +async def test_app() -> None: + async with create_async_test_client(asgi_app, middleware=[LoggingMiddlewareConfig().middleware]) as client: + response = await client.get("/") + assert response.status_code == 200 + assert response.text == "data: 1\ndata: 2\n"