Add max_event_size to cap SSE event buffering

docs/sse.md

@@ -74,4 +74,16 @@ If the response does not have a `text/event-stream` content type, iterating the
 
 `SSEError` is a subclass of [`TransportError`](exceptions.md), so it is also caught by `except httpx2.TransportError`.
 
+## Limiting event size
+
+An event is buffered until its terminating blank line arrives. To stop a stream that keeps sending data without ever completing an event from growing the buffer indefinitely, `client.sse()` caps the bytes buffered for a single event at 1 MiB by default and raises `SSEError` once an event exceeds the limit. Pass `max_event_size` to change the cap:
+
+```pycon
+>>> with client.sse("https://example.com/sse", max_event_size=8 * 1024 * 1024) as source:
+...     for event in source:  # raise the 1 MiB default to 8 MiB
+...         print(event.data)
+```
+
+The counter resets after each event, so the limit applies per event rather than to the stream as a whole. Set `max_event_size=None` to buffer events without any limit.
+
 [mdn]: https://developer.mozilla.org/en-US/docs/Web/API/Server-sent_events

src/httpx2/httpx2/_client.py

@@ -16,6 +16,7 @@
     DEFAULT_KEEPALIVE_PING_INTERVAL_SECONDS,
     DEFAULT_KEEPALIVE_PING_TIMEOUT_SECONDS,
     DEFAULT_LIMITS,
+    DEFAULT_MAX_EVENT_SIZE_BYTES,
     DEFAULT_MAX_MESSAGE_SIZE_BYTES,
     DEFAULT_MAX_REDIRECTS,
     DEFAULT_QUEUE_SIZE,
@@ -871,6 +872,7 @@ def sse(
         follow_redirects: bool | UseClientDefault = USE_CLIENT_DEFAULT,
         timeout: TimeoutTypes | UseClientDefault = USE_CLIENT_DEFAULT,
         extensions: RequestExtensions | None = None,
+        max_event_size: int | None = DEFAULT_MAX_EVENT_SIZE_BYTES,
     ) -> Generator[EventSource]:
         """
         Connect to a server-sent events endpoint and yield an `EventSource`.
@@ -894,7 +896,7 @@ def sse(
             timeout=timeout,
             extensions=extensions,
         ) as response:
-            yield EventSource(response)
+            yield EventSource(response, max_event_size=max_event_size)
 
     @contextmanager
     def websocket(
@@ -1708,6 +1710,7 @@ async def sse(
         follow_redirects: bool | UseClientDefault = USE_CLIENT_DEFAULT,
         timeout: TimeoutTypes | UseClientDefault = USE_CLIENT_DEFAULT,
         extensions: RequestExtensions | None = None,
+        max_event_size: int | None = DEFAULT_MAX_EVENT_SIZE_BYTES,
     ) -> AsyncGenerator[EventSource]:
         """
         Connect to a server-sent events endpoint and yield an `EventSource`.
@@ -1731,7 +1734,7 @@ async def sse(
             timeout=timeout,
             extensions=extensions,
         ) as response:
-            yield EventSource(response)
+            yield EventSource(response, max_event_size=max_event_size)
 
     @asynccontextmanager
     async def websocket(

src/httpx2/httpx2/_config.py

@@ -248,6 +248,8 @@ def __repr__(self) -> str:
 DEFAULT_LIMITS = Limits(max_connections=100, max_keepalive_connections=20)
 DEFAULT_MAX_REDIRECTS = 20
 
+DEFAULT_MAX_EVENT_SIZE_BYTES = 1024 * 1024
+
 DEFAULT_MAX_MESSAGE_SIZE_BYTES = 65_536
 DEFAULT_QUEUE_SIZE = 512
 DEFAULT_KEEPALIVE_PING_INTERVAL_SECONDS = 20.0

src/httpx2/httpx2/_sse.py

@@ -10,7 +10,8 @@
 from collections.abc import AsyncIterator, Iterator
 from dataclasses import dataclass
 
-from ._exceptions import TransportError
+from ._config import DEFAULT_MAX_EVENT_SIZE_BYTES
+from ._exceptions import TransportError, request_context
 from ._models import Response
 
 __all__ = ["EventSource", "SSEError", "ServerSentEvent"]
@@ -34,15 +35,18 @@ def json(self) -> object:
 
 
 class _SSEDecoder:
-    def __init__(self) -> None:
+    def __init__(self, max_event_size: int | None = None) -> None:
+        self._max_event_size = max_event_size
         self._event = ""
         self._data: list[str] = []
+        self._event_size = 0
         self._last_event_id = ""
         self._retry: int | None = None
         self._pending = False
 
     def decode(self, line: str) -> ServerSentEvent | None:
         if not line:
+            self._event_size = 0
             if not self._pending:
                 return None
 
@@ -61,6 +65,9 @@ def decode(self, line: str) -> ServerSentEvent | None:
         if line.startswith(":"):
             return None
 
+        self._event_size += len(line.encode("utf-8"))
+        self._check_size()
+
         fieldname, _, value = line.partition(":")
         value = value[1:] if value.startswith(" ") else value
 
@@ -83,12 +90,27 @@ def decode(self, line: str) -> ServerSentEvent | None:
 
         return None
 
+    def check_pending(self, pending: str) -> None:
+        """
+        Bound the total bytes buffered for the in-progress event, including
+        a trailing line that has not been terminated by a newline yet.
+        """
+        self._check_size(len(pending.encode("utf-8")))
+
+    def _check_size(self, pending_size: int = 0) -> None:
+        if self._max_event_size is not None and self._event_size + pending_size > self._max_event_size:
+            raise SSEError(f"Server-sent event exceeded the {self._max_event_size} byte limit.")
+
 
 class _SSELineDecoder:
     def __init__(self) -> None:
         self._buffer = ""
         self._trailing_cr = False
 
+    @property
+    def pending(self) -> str:
+        return self._buffer
+
     def decode(self, text: str) -> list[str]:
         if self._trailing_cr:
             text = "\r" + text
@@ -114,8 +136,9 @@ def flush(self) -> list[str]:
 
 
 class EventSource:
-    def __init__(self, response: Response) -> None:
+    def __init__(self, response: Response, max_event_size: int | None = DEFAULT_MAX_EVENT_SIZE_BYTES) -> None:
         self._response = response
+        self._max_event_size = max_event_size
 
     @property
     def response(self) -> Response:
@@ -124,35 +147,36 @@ def response(self) -> Response:
     def _check_content_type(self) -> None:
         content_type, _, _ = self._response.headers.get("content-type", "").partition(";")
         if content_type.strip().lower() != "text/event-stream":
-            raise SSEError(
-                f"Expected response with content type 'text/event-stream', got {content_type.strip()!r}.",
-                request=self._response.request,
-            )
+            raise SSEError(f"Expected response with content type 'text/event-stream', got {content_type.strip()!r}.")
 
     def __iter__(self) -> Iterator[ServerSentEvent]:
-        self._check_content_type()
-        decoder = _SSEDecoder()
-        lines = _SSELineDecoder()
-        for chunk in self._response.iter_text():
-            for line in lines.decode(chunk):
+        with request_context(request=self._response.request):
+            self._check_content_type()
+            decoder = _SSEDecoder(self._max_event_size)
+            lines = _SSELineDecoder()
+            for chunk in self._response.iter_text():
+                for line in lines.decode(chunk):
+                    sse = decoder.decode(line)
+                    if sse is not None:
+                        yield sse
+                decoder.check_pending(lines.pending)
+            for line in lines.flush():
                 sse = decoder.decode(line)
                 if sse is not None:
                     yield sse
-        for line in lines.flush():
-            sse = decoder.decode(line)
-            if sse is not None:
-                yield sse
 
     async def __aiter__(self) -> AsyncIterator[ServerSentEvent]:
-        self._check_content_type()
-        decoder = _SSEDecoder()
-        lines = _SSELineDecoder()
-        async for chunk in self._response.aiter_text():
-            for line in lines.decode(chunk):
+        with request_context(request=self._response.request):
+            self._check_content_type()
+            decoder = _SSEDecoder(self._max_event_size)
+            lines = _SSELineDecoder()
+            async for chunk in self._response.aiter_text():
+                for line in lines.decode(chunk):
+                    sse = decoder.decode(line)
+                    if sse is not None:
+                        yield sse
+                decoder.check_pending(lines.pending)
+            for line in lines.flush():
                 sse = decoder.decode(line)
                 if sse is not None:
                     yield sse
-        for line in lines.flush():
-            sse = decoder.decode(line)
-            if sse is not None:
-                yield sse