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,10 @@ def decode(self, line: str) -> ServerSentEvent | None:
         if line.startswith(":"):
             return None
 
+        self._event_size += len(line.encode("utf-8"))
+        if self._max_event_size is not None and self._event_size > self._max_event_size:
+            raise SSEError(f"Server-sent event exceeded the {self._max_event_size} byte limit.")
+
         fieldname, _, value = line.partition(":")
         value = value[1:] if value.startswith(" ") else value
 
@@ -85,7 +93,8 @@ def decode(self, line: str) -> ServerSentEvent | None:
 
 
 class _SSELineDecoder:
-    def __init__(self) -> None:
+    def __init__(self, max_event_size: int | None = None) -> None:
+        self._max_event_size = max_event_size
         self._buffer = ""
         self._trailing_cr = False
 
@@ -100,6 +109,8 @@ def decode(self, text: str) -> list[str]:
         text = self._buffer + text.replace("\r\n", "\n").replace("\r", "\n")
         lines = text.split("\n")
         self._buffer = lines.pop()
+        if self._max_event_size is not None and len(self._buffer.encode("utf-8")) > self._max_event_size:
+            raise SSEError(f"Server-sent event exceeded the {self._max_event_size} byte limit.")
         return lines
 
     def flush(self) -> list[str]:
@@ -114,8 +125,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 +136,34 @@ 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(self._max_event_size)
+            for chunk in self._response.iter_text():
+                for line in lines.decode(chunk):
+                    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
-        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(self._max_event_size)
+            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
+            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

tests/httpx2/test_sse.py

@@ -275,6 +275,199 @@ def handler(request: httpx2.Request) -> httpx2.Response:
     assert event.data == "hi"
 
 
+def test_max_event_size_rejects_unterminated_line() -> None:
+    def handler(request: httpx2.Request) -> httpx2.Response:
+        body = b"data: " + b"A" * 8192
+        return httpx2.Response(200, content=body, headers={"Content-Type": "text/event-stream"})
+
+    with httpx2.Client(transport=httpx2.MockTransport(handler)) as client:
+        with client.sse("http://testserver/sse", max_event_size=4096) as source:
+            with pytest.raises(httpx2.SSEError, match="4096 byte limit"):
+                list(source)
+
+
+def test_max_event_size_rejects_many_data_lines() -> None:
+    def handler(request: httpx2.Request) -> httpx2.Response:
+        body = b"data: chunk\n" * 1000
+        return httpx2.Response(200, content=body, headers={"Content-Type": "text/event-stream"})
+
+    with httpx2.Client(transport=httpx2.MockTransport(handler)) as client:
+        with client.sse("http://testserver/sse", max_event_size=100) as source:
+            with pytest.raises(httpx2.SSEError, match="100 byte limit"):
+                list(source)
+
+
+def test_max_event_size_counts_empty_data_lines() -> None:
+    def handler(request: httpx2.Request) -> httpx2.Response:
+        body = b"data\n" * 1000
+        return httpx2.Response(200, content=body, headers={"Content-Type": "text/event-stream"})
+
+    with httpx2.Client(transport=httpx2.MockTransport(handler)) as client:
+        with client.sse("http://testserver/sse", max_event_size=100) as source:
+            with pytest.raises(httpx2.SSEError, match="100 byte limit"):
+                list(source)
+
+
+def test_max_event_size_counts_non_data_fields() -> None:
+    def handler(request: httpx2.Request) -> httpx2.Response:
+        body = b"id: " + b"A" * 8192 + b"\n\n"
+        return httpx2.Response(200, content=body, headers={"Content-Type": "text/event-stream"})
+
+    with httpx2.Client(transport=httpx2.MockTransport(handler)) as client:
+        with client.sse("http://testserver/sse", max_event_size=4096) as source:
+            with pytest.raises(httpx2.SSEError, match="4096 byte limit"):
+                list(source)
+
+
+def test_max_event_size_measures_utf8_bytes() -> None:
+    def handler(request: httpx2.Request) -> httpx2.Response:
+        body = b"data: " + "😀".encode() * 50 + b"\n\n"
+        return httpx2.Response(200, content=body, headers={"Content-Type": "text/event-stream"})
+
+    with httpx2.Client(transport=httpx2.MockTransport(handler)) as client:
+        with client.sse("http://testserver/sse", max_event_size=100) as source:
+            with pytest.raises(httpx2.SSEError, match="100 byte limit"):
+                list(source)
+
+
+def test_max_event_size_rejects_single_large_data_line() -> None:
+    def handler(request: httpx2.Request) -> httpx2.Response:
+        body = b"data: " + b"A" * 8192 + b"\n\n"
+        return httpx2.Response(200, content=body, headers={"Content-Type": "text/event-stream"})
+
+    with httpx2.Client(transport=httpx2.MockTransport(handler)) as client:
+        with client.sse("http://testserver/sse", max_event_size=4096) as source:
+            with pytest.raises(httpx2.SSEError, match="4096 byte limit"):
+                list(source)
+
+
+def test_max_event_size_ignores_keepalive_comments() -> None:
+    def handler(request: httpx2.Request) -> httpx2.Response:
+        body = b": keepalive\n\n" * 1000 + b"data: hi\n\n"
+        return httpx2.Response(200, content=body, headers={"Content-Type": "text/event-stream"})
+
+    with httpx2.Client(transport=httpx2.MockTransport(handler)) as client:
+        with client.sse("http://testserver/sse", max_event_size=100) as source:
+            (event,) = list(source)
+
+    assert event.data == "hi"
+
+
+def test_max_event_size_allows_event_under_limit() -> None:
+    def handler(request: httpx2.Request) -> httpx2.Response:
+        return httpx2.Response(200, content=b"data: hi\n\n", headers={"Content-Type": "text/event-stream"})
+
+    with httpx2.Client(transport=httpx2.MockTransport(handler)) as client:
+        with client.sse("http://testserver/sse", max_event_size=4096) as source:
+            (event,) = list(source)
+
+    assert event.data == "hi"
+
+
+def test_max_event_size_resets_between_events() -> None:
+    def handler(request: httpx2.Request) -> httpx2.Response:
+        body = b"data: " + b"A" * 80 + b"\n\ndata: " + b"B" * 80 + b"\n\n"
+        return httpx2.Response(200, content=body, headers={"Content-Type": "text/event-stream"})
+
+    with httpx2.Client(transport=httpx2.MockTransport(handler)) as client:
+        with client.sse("http://testserver/sse", max_event_size=100) as source:
+            events = list(source)
+
+    assert [len(event.data) for event in events] == [80, 80]
+
+
+def test_max_event_size_applies_by_default() -> None:
+    def handler(request: httpx2.Request) -> httpx2.Response:
+        body = b"data: " + b"A" * (2 * 1024 * 1024) + b"\n\n"
+        return httpx2.Response(200, content=body, headers={"Content-Type": "text/event-stream"})
+
+    with httpx2.Client(transport=httpx2.MockTransport(handler)) as client:
+        with client.sse("http://testserver/sse") as source:
+            with pytest.raises(httpx2.SSEError, match="byte limit"):
+                list(source)
+
+
+def test_max_event_size_none_disables_limit() -> None:
+    def handler(request: httpx2.Request) -> httpx2.Response:
+        body = b"data: " + b"A" * (2 * 1024 * 1024) + b"\n\n"
+        return httpx2.Response(200, content=body, headers={"Content-Type": "text/event-stream"})
+
+    with httpx2.Client(transport=httpx2.MockTransport(handler)) as client:
+        with client.sse("http://testserver/sse", max_event_size=None) as source:
+            (event,) = list(source)
+
+    assert len(event.data) == 2 * 1024 * 1024
+
+
+def test_max_event_size_spans_completed_and_pending_lines() -> None:
+    def chunks() -> Iterator[bytes]:
+        yield b"data: " + b"A" * 60 + b"\n"
+        yield b"data: " + b"B" * 60
+
+    def handler(request: httpx2.Request) -> httpx2.Response:
+        return httpx2.Response(200, content=chunks(), headers={"Content-Type": "text/event-stream"})
+
+    with httpx2.Client(transport=httpx2.MockTransport(handler)) as client:
+        with client.sse("http://testserver/sse", max_event_size=100) as source:
+            with pytest.raises(httpx2.SSEError, match="100 byte limit"):
+                list(source)
+
+
+def test_max_event_size_does_not_conflate_adjacent_events() -> None:
+    def chunks() -> Iterator[bytes]:
+        yield b"data: " + b"A" * 60 + b"\n\ndata: " + b"B" * 60
+        yield b"\n\n"
+
+    def handler(request: httpx2.Request) -> httpx2.Response:
+        return httpx2.Response(200, content=chunks(), headers={"Content-Type": "text/event-stream"})
+
+    with httpx2.Client(transport=httpx2.MockTransport(handler)) as client:
+        with client.sse("http://testserver/sse", max_event_size=100) as source:
+            events = list(source)
+
+    assert [len(event.data) for event in events] == [60, 60]
+
+
+def test_max_event_size_error_has_request() -> None:
+    def handler(request: httpx2.Request) -> httpx2.Response:
+        body = b"data: " + b"A" * 8192 + b"\n\n"
+        return httpx2.Response(200, content=body, headers={"Content-Type": "text/event-stream"})
+
+    with httpx2.Client(transport=httpx2.MockTransport(handler)) as client:
+        with client.sse("http://testserver/sse", max_event_size=4096) as source:
+            with pytest.raises(httpx2.SSEError) as exc_info:
+                list(source)
+
+    assert exc_info.value.request.url == "http://testserver/sse"
+
+
+@pytest.mark.anyio
+async def test_max_event_size_applies_by_default_async() -> None:
+    captured: list[int | None] = []
+
+    def handler(request: httpx2.Request) -> httpx2.Response:
+        return httpx2.Response(200, content=b"data: hi\n\n", headers={"Content-Type": "text/event-stream"})
+
+    async with httpx2.AsyncClient(transport=httpx2.MockTransport(handler)) as client:
+        async with client.sse("http://testserver/sse") as source:
+            captured.append(source._max_event_size)
+            [event async for event in source]
+
+    assert captured == [1024 * 1024]
+
+
+@pytest.mark.anyio
+async def test_max_event_size_allows_event_under_limit_async() -> None:
+    def handler(request: httpx2.Request) -> httpx2.Response:
+        return httpx2.Response(200, content=b"data: hi\n\n", headers={"Content-Type": "text/event-stream"})
+
+    async with httpx2.AsyncClient(transport=httpx2.MockTransport(handler)) as client:
+        async with client.sse("http://testserver/sse", max_event_size=4096) as source:
+            events = [event async for event in source]
+
+    assert [event.data for event in events] == ["hi"]
+
+
 def test_event_dispatched_at_eof_on_trailing_cr_sync() -> None:
     def chunks() -> Iterator[bytes]:
         yield b"data: hi\n"