diff --git a/src/httpx2/httpx2/_sse.py b/src/httpx2/httpx2/_sse.py index 9834baea..df2be6aa 100644 --- a/src/httpx2/httpx2/_sse.py +++ b/src/httpx2/httpx2/_sse.py @@ -112,12 +112,12 @@ def decode(self, line: str) -> ServerSentEvent | None: return None - def check_pending(self, pending: str) -> None: + def check_pending(self, pending_size: int) -> 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"))) + self._check_size(pending_size) 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: @@ -126,12 +126,13 @@ def _check_size(self, pending_size: int = 0) -> None: class _SSELineDecoder: def __init__(self) -> None: - self._buffer = "" + self._parts: list[str] = [] + self._pending_size = 0 self._trailing_cr = False @property - def pending(self) -> str: - return self._buffer + def pending_size(self) -> int: + return self._pending_size def decode(self, text: str) -> list[str]: if self._trailing_cr: @@ -141,20 +142,36 @@ def decode(self, text: str) -> list[str]: self._trailing_cr = True text = text[:-1] - text = self._buffer + text.replace("\r\n", "\n").replace("\r", "\n") + text = text.replace("\r\n", "\n").replace("\r", "\n") + if "\n" not in text: + self._append(text) + return [] + lines = text.split("\n") - self._buffer = lines.pop() + self._append(lines[0]) + lines[0] = self._consume_pending() + self._append(lines.pop()) return lines def flush(self) -> list[str]: if self._trailing_cr: - self._buffer += "\n" + self._append("\n") self._trailing_cr = False - if not self._buffer: + buffer = self._consume_pending() + if not buffer: return [] - lines = self._buffer.split("\n") - self._buffer = "" - return lines + return buffer.split("\n") + + def _append(self, text: str) -> None: + if text: + self._parts.append(text) + self._pending_size += len(text.encode("utf-8")) + + def _consume_pending(self) -> str: + pending = "".join(self._parts) + self._parts = [] + self._pending_size = 0 + return pending class _SSEParser: @@ -164,7 +181,7 @@ def __init__(self, max_event_size: int | None = None) -> None: def decode(self, text: str) -> Iterator[ServerSentEvent]: yield from self._decode_lines(self._line_decoder.decode(text)) - self._event_decoder.check_pending(self._line_decoder.pending) + self._event_decoder.check_pending(self._line_decoder.pending_size) def flush(self) -> Iterator[ServerSentEvent]: yield from self._decode_lines(self._line_decoder.flush()) diff --git a/tests/httpx2/test_sse.py b/tests/httpx2/test_sse.py index fb15e63b..7aaad5b0 100644 --- a/tests/httpx2/test_sse.py +++ b/tests/httpx2/test_sse.py @@ -497,3 +497,21 @@ def handler(request: httpx2.Request) -> httpx2.Response: events = [event async for event in source] assert [event.data for event in events] == ["hi"] + + +def test_many_chunks_without_line_separator() -> None: + parts = [f"{index:04d}".encode() for index in range(1_000)] + + def chunks() -> Iterator[bytes]: + yield b"data: " + yield from parts + 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") as source: + (event,) = list(source) + + assert event.data == b"".join(parts).decode()