diff --git a/baseten/sandbox/_common.py b/baseten/sandbox/_common.py index fba8107..c95c228 100644 --- a/baseten/sandbox/_common.py +++ b/baseten/sandbox/_common.py @@ -437,9 +437,12 @@ def response_lines(response: httpx.Response) -> Generator[str, None, None]: decoder = codecs.getincrementaldecoder("utf-8")(errors="replace") buffer = "" for chunk in response.iter_bytes(): - buffer += decoder.decode(chunk) - *lines, buffer = _LINE_BREAK.split(buffer) - yield from lines + text = decoder.decode(chunk) + buffer += text + # Split only when a line ends, so a long line isn't rescanned per chunk. + if "\n" in text: + *lines, buffer = _LINE_BREAK.split(buffer) + yield from lines buffer += decoder.decode(b"", final=True) if buffer: yield buffer @@ -450,10 +453,12 @@ async def response_lines_async(response: httpx.Response) -> AsyncGenerator[str, decoder = codecs.getincrementaldecoder("utf-8")(errors="replace") buffer = "" async for chunk in response.aiter_bytes(): - buffer += decoder.decode(chunk) - *lines, buffer = _LINE_BREAK.split(buffer) - for line in lines: - yield line + text = decoder.decode(chunk) + buffer += text + if "\n" in text: + *lines, buffer = _LINE_BREAK.split(buffer) + for line in lines: + yield line buffer += decoder.decode(b"", final=True) if buffer: yield buffer diff --git a/tests/sandbox/test_common.py b/tests/sandbox/test_common.py index b8e4949..cb9f460 100644 --- a/tests/sandbox/test_common.py +++ b/tests/sandbox/test_common.py @@ -146,6 +146,8 @@ async def aclose(self) -> None: # A character split across chunks, and invalid UTF-8 replaced. (("é\n".encode()[:1], "é\n".encode()[1:]), ["é"]), ((b"a\xffb\n",), ["a�b"]), + # A long line across many chunks, ended by \r\n split across two. + ((*[b"x" * 1024] * 64, b"\r", b"\ny"), ["x" * 65536, "y"]), ], ) class TestResponseLines: