fix: logging middleware with multi-body response (#3478)

* fix: logging middleware with multi-body response

Prevent logging middleware from failing with a KeyError when a response sends multiple "http.response.body" messages.

Closes #3477

* fix: ensure asgi messages aren't kept in log context

Author: peterschutt

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

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"