Repository: litestar-org/litestar

Issue:
Bug: Multi-body response incompatible with LoggingMiddleware
### Description

When using `ServerSentEvent` Responses with `StructlogPlugin`, the application errors. 

Preliminary research led me to https://github.com/litestar-org/litestar/blob/main/litestar/middleware/logging.py#L180, where `scope.state._ls_connection_state.log_context` does not have any values during second push of an `ServerSentEventMessage`. 

### URL to code causing the issue

_No response_

### MCVE

```python
from asyncio import sleep
from collections.abc import AsyncGenerator

from litestar import Litestar, get
from litestar.plugins.structlog import StructlogPlugin
from litestar.response import ServerSentEvent, ServerSentEventMessage
from litestar.types import SSEData


async def my_generator() -> AsyncGenerator[SSEData, None]:
    count = 0
    while count < 10:
        await sleep(0.01)
        count += 1
        yield ServerSentEventMessage(event="something-with-comment", retry=1000, comment="some comment")


@get(path="/count", sync_to_thread=False)
def sse_handler() -> ServerSentEvent:
    return ServerSentEvent(my_generator())


app = Litestar(route_handlers=[sse_handler],
               plugins=[StructlogPlugin()])
```


### Steps to reproduce

1. Run the server above

```bash
litestar --app main:app run
```

2. Send a request

```bash
 curl localhost:8000/count
```

3. Observe error

### Screenshots

_No response_

### Logs

```bash
File "/.venv/lib/python3.11/site-packages/litestar/middleware/logging.py", line 194, in extract_response_data
     connection_state.log_context.pop(HTTP_RESPONSE_START),
     ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
KeyError: 'http.response.start'
```


### Litestar Version

2.8.3

### Platform

- [ ] Linux
- [X] Mac
- [ ] Windows
- [ ] Other (Please specify in the description above)

<!-- POLAR PLEDGE BADGE START -->
---
> [!NOTE]  
> While we are open for sponsoring on [GitHub Sponsors](https://github.com/sponsors/litestar-org/) and 
> [OpenCollective](https://opencollective.com/litestar), we also utilize [Polar.sh](https://polar.sh/) to engage in pledge-based sponsorship.
>
> Check out all issues funded or available for funding [on our Polar.sh dashboard](https://polar.sh/litestar-org)
> * If you would like to see an issue prioritized, make a pledge towards it!
> * We receive the pledge once the issue is completed & verified
> * This, along with engagement in the community, helps us know which features are a priority to our users.

<a href="https://polar.sh/litestar-org/litestar/issues/3477">
<picture>
  <source media="(prefers-color-scheme: dark)" srcset="https://polar.sh/api/github/litestar-org/litestar/issues/3477/pledge.svg?darkmode=1">
  <img alt="Fund with Polar" src="https://polar.sh/api/github/litestar-org/litestar/issues/3477/pledge.svg">
</picture>
</a>
<!-- POLAR PLEDGE BADGE END -->



Discussion:
This generalizes down to an incompatibility between logging middleware and any multi-body response:

```py
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"
```

Resolving patch:
diff --git a/litestar/middleware/logging.py b/litestar/middleware/logging.py
--- 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


Test patch:
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
--- /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"