From 69516b1d5c96888c5a42738ed831c2705b677617 Mon Sep 17 00:00:00 2001 From: Ayush Agrawal Date: Thu, 29 Oct 2026 02:18:14 -0700 Subject: [PATCH] fix: Split SSE stream lines on CR, LF and CRLF only fixes #3243 FUTURE_COPYBARA_INTEGRATE_REVIEW=https://github.com/googleapis/python-genai/pull/3209 from googleapis:release-please--branches--main 39c95f2f004970ff5d9fd941aef3a09a98d3a79e PiperOrigin-RevId: 1009528460 --- google/genai/_api_client.py | 76 ++++++- .../genai/tests/client/test_async_stream.py | 212 ++++++++++++++++-- 2 files changed, 271 insertions(+), 17 deletions(-) diff --git a/google/genai/_api_client.py b/google/genai/_api_client.py index 36d7a2638d..b015d002a4 100644 --- a/google/genai/_api_client.py +++ b/google/genai/_api_client.py @@ -20,7 +20,9 @@ """ import asyncio +from collections.abc import AsyncIterable from collections.abc import Generator +from collections.abc import Iterable import copy from dataclasses import dataclass import inspect @@ -362,6 +364,71 @@ class HttpRequest: timeout: Optional[float] = None +class _SseLineDecoder: + """Splits a byte stream into lines at CR, LF or CRLF only. + + The SSE format allows no other line terminators. str.splitlines(), which + httpx and requests use, also breaks at U+2028, U+2029 and U+0085, and these + may appear unescaped inside JSON strings. + """ + + def __init__(self) -> None: + self._buffer: list[bytes] = [] + self._trailing_cr = False + + def decode(self, chunk: bytes) -> list[str]: + """Returns the lines that `chunk` completes.""" + if self._trailing_cr: + chunk = b'\r' + chunk + self._trailing_cr = False + # Hold a trailing CR until the next chunk shows whether LF follows. + if chunk.endswith(b'\r'): + self._trailing_cr = True + chunk = chunk[:-1] + if not chunk: + return [] + + # bytes.splitlines() breaks only at \r, \n and \r\n. + lines = chunk.splitlines() + tail = None if chunk.endswith((b'\r', b'\n')) else lines.pop() + if lines: + self._buffer.append(lines[0]) + lines[0] = b''.join(self._buffer) + self._buffer = [] + if tail is not None: + self._buffer.append(tail) + return [line.decode('utf-8', errors='replace') for line in lines] + + def flush(self) -> list[str]: + """Returns the buffered line at the end of the stream, if any.""" + if not self._buffer and not self._trailing_cr: + return [] + line = b''.join(self._buffer) + self._buffer = [] + self._trailing_cr = False + return [line.decode('utf-8', errors='replace')] + + +def _iter_sse_lines(chunks: Iterable[bytes]) -> Iterator[str]: + """Yields the lines of an SSE byte stream.""" + decoder = _SseLineDecoder() + for chunk in chunks: + yield from decoder.decode(chunk) + yield from decoder.flush() + + +async def _aiter_sse_lines( + chunks: AsyncIterable[bytes], +) -> AsyncIterator[str]: + """Yields the lines of an async SSE byte stream.""" + decoder = _SseLineDecoder() + async for chunk in chunks: + for line in decoder.decode(chunk): + yield line + for line in decoder.flush(): + yield line + + class HttpResponse: def __init__( @@ -485,9 +552,12 @@ def _iter_response_stream(self) -> Iterator[str]: balance = 0 data_buffer: list[str] = [] if isinstance(self.response_stream, _HTTPX_RESPONSE_TYPES): - response_stream = self.response_stream.iter_lines() + response_stream = _iter_sse_lines(self.response_stream.iter_bytes()) else: - response_stream = self.response_stream.iter_lines(decode_unicode=True) + # 512 bytes is the chunk size requests' iter_lines() reads. + response_stream = _iter_sse_lines( + self.response_stream.iter_content(chunk_size=512) + ) for line in response_stream: if not line: if data_buffer: @@ -542,7 +612,7 @@ async def _aiter_response_stream(self) -> AsyncIterator[str]: if isinstance(self.response_stream, _HTTPX_RESPONSE_TYPES): try: response_stream: Any = self.response_stream - async for line in response_stream.aiter_lines(): + async for line in _aiter_sse_lines(response_stream.aiter_bytes()): if not line: if data_buffer: yield '\n'.join(data_buffer) diff --git a/google/genai/tests/client/test_async_stream.py b/google/genai/tests/client/test_async_stream.py index 71b626b9f7..89970b323d 100644 --- a/google/genai/tests/client/test_async_stream.py +++ b/google/genai/tests/client/test_async_stream.py @@ -16,12 +16,15 @@ """Tests for async stream.""" import asyncio -from typing import List +import io +import json +from typing import Any, List from unittest import mock from unittest.mock import AsyncMock from unittest.mock import MagicMock from unittest.mock import patch import pytest +import requests try: @@ -51,15 +54,12 @@ class MockHTTPXResponse(httpx.Response): """Mock httpx.Response class for testing.""" def __init__(self, lines: List[str]): - self.aiter_lines = MagicMock() - self.aiter_lines.return_value.__aiter__ = MagicMock( - return_value=self._async_line_iterator(lines) - ) + self.aiter_bytes = MagicMock(return_value=self._async_byte_iterator(lines)) self.aclose = AsyncMock() - async def _async_line_iterator(self, lines: List[str]): + async def _async_byte_iterator(self, lines: List[str]): for line in lines: - yield line + yield line.encode('utf-8') + b'\n' class MockAIOHTTPResponse(aiohttp.ClientResponse): @@ -128,7 +128,7 @@ async def test_httpx_simple_lines(responses: api_client.HttpResponse): results = [line async for line in responses._aiter_response_stream()] assert results == lines - mock_response.aiter_lines.assert_called_once() + mock_response.aiter_bytes.assert_called_once() mock_response.aclose.assert_called_once() @@ -141,7 +141,7 @@ async def test_httpx_data_prefix(responses: api_client.HttpResponse): results = [line async for line in responses._aiter_response_stream()] assert results == ["{ 'message': 'hello' }", "{ 'status': 'ok' }"] - mock_response.aiter_lines.assert_called_once() + mock_response.aiter_bytes.assert_called_once() mock_response.aclose.assert_called_once() @pytest.mark.asyncio @@ -158,7 +158,7 @@ async def test_httpx_multiline_data(responses: api_client.HttpResponse) -> None: results = [line async for line in responses._aiter_response_stream()] assert results == ["{\n 'message': 'hello'\n}"] - mock_response.aiter_lines.assert_called_once() + mock_response.aiter_bytes.assert_called_once() mock_response.aclose.assert_called_once() @@ -170,7 +170,9 @@ def test_httpx_multiline_data_sync(responses: api_client.HttpResponse) -> None: "", ] mock_response = MagicMock(spec=httpx.Response) - mock_response.iter_lines.return_value = iter(lines) + mock_response.iter_bytes.return_value = iter( + [b'\n'.join(line.encode('utf-8') for line in lines) + b'\n'] + ) responses.response_stream = mock_response results = list(responses._iter_response_stream()) @@ -193,7 +195,7 @@ async def test_httpx_multiple_json_chunk(responses: api_client.HttpResponse): results = [line async for line in responses._aiter_response_stream()] assert results == ['{ "id": 1 }', '{ "id": 2 }', '{ "id": 3 }'] - mock_response.aiter_lines.assert_called_once() + mock_response.aiter_bytes.assert_called_once() mock_response.aclose.assert_called_once() @@ -207,7 +209,7 @@ async def test_httpx_incomplete_json_at_end(responses: api_client.HttpResponse): # The remaining chunk is yielded assert results == ['{ "partial": "data"'] - mock_response.aiter_lines.assert_called_once() + mock_response.aiter_bytes.assert_called_once() mock_response.aclose.assert_called_once() @@ -220,7 +222,7 @@ async def test_httpx_empty_stream(responses: api_client.HttpResponse): results = [line async for line in responses._aiter_response_stream()] assert results == [] - mock_response.aiter_lines.assert_called_once() + mock_response.aiter_bytes.assert_called_once() mock_response.aclose.assert_called_once() @@ -538,3 +540,185 @@ async def test_error_event_in_generate_content_stream_async(mock_transport): with pytest.raises(errors.ServerError): await anext(generated_response) + + +_LINE_SEPARATORS = pytest.mark.parametrize( + 'sep', ['\u2028', '\u2029', '\x85'], ids=['u2028', 'u2029', 'u0085'] +) + + +def _sse_body_with_separator(sep: str) -> bytes: + return b''.join( + b'data: ' + + json.dumps(c, ensure_ascii=False).encode('utf-8') + + b'\r\n\r\n' + for c in _sse_events_with_separator(sep) + ) + + +def _sse_events_with_separator(sep: str) -> list[dict[str, Any]]: + return [ + { + 'candidates': [{ + 'content': { + 'role': 'model', + 'parts': [{'text': f'line one{sep}line two'}], + } + }] + }, + { + 'candidates': [{ + 'content': {'role': 'model', 'parts': [{'text': ' done'}]}, + 'finishReason': 'STOP', + }] + }, + ] + + +def _sse_handler(sep: str): + def handler(_: httpx.Request) -> httpx.Response: + return httpx.Response( + 200, + content=_sse_body_with_separator(sep), + headers={'content-type': 'text/event-stream'}, + ) + + return handler + + +@_LINE_SEPARATORS +def test_generate_content_stream_line_separator_in_text(sep): + """Unicode line separators inside a JSON string do not split an SSE line.""" + client = Client( + api_key='test_api_key', + http_options=types.HttpOptions( + client_args={'transport': httpx.MockTransport(_sse_handler(sep))} + ), + ) + + stream = client.models.generate_content_stream( + model='gemini-2.5-flash', contents='hi' + ) + + assert [c.text for c in stream] == [f'line one{sep}line two', ' done'] + + +@_LINE_SEPARATORS +@pytest.mark.asyncio +async def test_generate_content_stream_line_separator_in_text_async(sep): + """Async httpx streams keep Unicode line separators inside JSON strings.""" + client = Client( + api_key='test_api_key', + http_options=types.HttpOptions( + async_client_args={ + 'transport': httpx.MockTransport(_sse_handler(sep)) + } + ), + ) + + stream = await client.aio.models.generate_content_stream( + model='gemini-2.5-flash', contents='hi' + ) + + assert [c.text async for c in stream] == [ + f'line one{sep}line two', + ' done', + ] + + +@_LINE_SEPARATORS +def test_requests_stream_line_separator_in_text(sep): + """The requests transport splits SSE lines only at CR, LF and CRLF.""" + resp = requests.Response() + resp.status_code = 200 + resp.raw = io.BytesIO(_sse_body_with_separator(sep)) + resp.encoding = 'utf-8' + + results = list( + api_client.HttpResponse( + headers={}, response_stream=resp + )._iter_response_stream() + ) + + assert results == [ + json.dumps(c, ensure_ascii=False) for c in _sse_events_with_separator(sep) + ] + + +def test_httpx_stream_crlf_split_across_chunks(): + """A CRLF split across two chunks ends exactly one line.""" + chunks = [b'data: {"a": 1}\r', b'\n\r\n', b'data: {"b": 2}\r\n\r\n'] + resp = api_client.HttpResponse( + headers={}, response_stream=mock_response(chunks) + ) + + assert list(resp._iter_response_stream()) == ['{"a": 1}', '{"b": 2}'] + + +def test_httpx_stream_cr_only_terminators(): + """A bare CR ends a line, as the SSE format requires.""" + chunks = [b'data: {"a": 1}\r\rdata: {"b": 2}\r\r'] + resp = api_client.HttpResponse( + headers={}, response_stream=mock_response(chunks) + ) + + assert list(resp._iter_response_stream()) == ['{"a": 1}', '{"b": 2}'] + + +def test_httpx_stream_multibyte_char_split_across_chunks(): + """A UTF-8 character split across chunks decodes as one character.""" + chunks = [b'data: {"t": "x\xe2\x80', b'\xa8y"}\n\n'] + resp = api_client.HttpResponse( + headers={}, response_stream=mock_response(chunks) + ) + + assert list(resp._iter_response_stream()) == ['{"t": "x\u2028y"}'] + + +def test_httpx_stream_multiline_data_with_line_separator(): + """Multi-line data fields still join with LF around a U+2028.""" + chunks = [b'data: {"t": "a\xe2\x80\xa8b",\n', b'data: "u": 1}\n\n'] + resp = api_client.HttpResponse( + headers={}, response_stream=mock_response(chunks) + ) + + assert list(resp._iter_response_stream()) == ['{"t": "a\u2028b",\n"u": 1}'] + + +def test_httpx_stream_invalid_utf8_is_replaced(): + """Invalid UTF-8 bytes become U+FFFD instead of raising.""" + chunks = [b'data: {"t": "\xff"}\n\n'] + resp = api_client.HttpResponse( + headers={}, response_stream=mock_response(chunks) + ) + + assert list(resp._iter_response_stream()) == ['{"t": "\ufffd"}'] + + +def test_httpx_stream_final_line_without_terminator(): + """A last line with no terminator is still yielded at end of stream.""" + chunks = [b'data: {"a": 1}\n\n{"error": ', b'{"code": 500}}'] + resp = api_client.HttpResponse( + headers={}, response_stream=mock_response(chunks) + ) + + assert list(resp._iter_response_stream()) == [ + '{"a": 1}', + '{"error": {"code": 500}}', + ] + + +@pytest.mark.asyncio +async def test_httpx_async_stream_line_separator_in_text( + responses: api_client.HttpResponse, +): + """The async httpx branch keeps a U+2028 line whole and closes once.""" + httpx_response = MockHTTPXResponse( + ['data: {"t": "a\u2028b"}', '', 'data: {"t": "c"}', ''] + ) + responses.response_stream = httpx_response + + results = [line async for line in responses._aiter_response_stream()] + + assert results == ['{"t": "a\u2028b"}', '{"t": "c"}'] + httpx_response.aclose.assert_called_once()