Skip to content

Commit 1665a33

Browse files
speedstorm1copybara-github
authored andcommitted
fix: Split SSE stream lines on CR, LF and CRLF only
fixes #3243 FUTURE_COPYBARA_INTEGRATE_REVIEW=#3209 from googleapis:release-please--branches--main 39c95f2 PiperOrigin-RevId: 1009528460
1 parent 0dc788c commit 1665a33

2 files changed

Lines changed: 271 additions & 17 deletions

File tree

‎google/genai/_api_client.py‎

Lines changed: 73 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -20,7 +20,9 @@
2020
"""
2121

2222
import asyncio
23+
from collections.abc import AsyncIterable
2324
from collections.abc import Generator
25+
from collections.abc import Iterable
2426
import copy
2527
from dataclasses import dataclass
2628
import inspect
@@ -362,6 +364,71 @@ class HttpRequest:
362364
timeout: Optional[float] = None
363365

364366

367+
class _SseLineDecoder:
368+
"""Splits a byte stream into lines at CR, LF or CRLF only.
369+
370+
The SSE format allows no other line terminators. str.splitlines(), which
371+
httpx and requests use, also breaks at U+2028, U+2029 and U+0085, and these
372+
may appear unescaped inside JSON strings.
373+
"""
374+
375+
def __init__(self) -> None:
376+
self._buffer: list[bytes] = []
377+
self._trailing_cr = False
378+
379+
def decode(self, chunk: bytes) -> list[str]:
380+
"""Returns the lines that `chunk` completes."""
381+
if self._trailing_cr:
382+
chunk = b'\r' + chunk
383+
self._trailing_cr = False
384+
# Hold a trailing CR until the next chunk shows whether LF follows.
385+
if chunk.endswith(b'\r'):
386+
self._trailing_cr = True
387+
chunk = chunk[:-1]
388+
if not chunk:
389+
return []
390+
391+
# bytes.splitlines() breaks only at \r, \n and \r\n.
392+
lines = chunk.splitlines()
393+
tail = None if chunk.endswith((b'\r', b'\n')) else lines.pop()
394+
if lines:
395+
self._buffer.append(lines[0])
396+
lines[0] = b''.join(self._buffer)
397+
self._buffer = []
398+
if tail is not None:
399+
self._buffer.append(tail)
400+
return [line.decode('utf-8', errors='replace') for line in lines]
401+
402+
def flush(self) -> list[str]:
403+
"""Returns the buffered line at the end of the stream, if any."""
404+
if not self._buffer and not self._trailing_cr:
405+
return []
406+
line = b''.join(self._buffer)
407+
self._buffer = []
408+
self._trailing_cr = False
409+
return [line.decode('utf-8', errors='replace')]
410+
411+
412+
def _iter_sse_lines(chunks: Iterable[bytes]) -> Iterator[str]:
413+
"""Yields the lines of an SSE byte stream."""
414+
decoder = _SseLineDecoder()
415+
for chunk in chunks:
416+
yield from decoder.decode(chunk)
417+
yield from decoder.flush()
418+
419+
420+
async def _aiter_sse_lines(
421+
chunks: AsyncIterable[bytes],
422+
) -> AsyncIterator[str]:
423+
"""Yields the lines of an async SSE byte stream."""
424+
decoder = _SseLineDecoder()
425+
async for chunk in chunks:
426+
for line in decoder.decode(chunk):
427+
yield line
428+
for line in decoder.flush():
429+
yield line
430+
431+
365432
class HttpResponse:
366433

367434
def __init__(
@@ -485,9 +552,12 @@ def _iter_response_stream(self) -> Iterator[str]:
485552
balance = 0
486553
data_buffer: list[str] = []
487554
if isinstance(self.response_stream, _HTTPX_RESPONSE_TYPES):
488-
response_stream = self.response_stream.iter_lines()
555+
response_stream = _iter_sse_lines(self.response_stream.iter_bytes())
489556
else:
490-
response_stream = self.response_stream.iter_lines(decode_unicode=True)
557+
# 512 bytes is the chunk size requests' iter_lines() reads.
558+
response_stream = _iter_sse_lines(
559+
self.response_stream.iter_content(chunk_size=512)
560+
)
491561
for line in response_stream:
492562
if not line:
493563
if data_buffer:
@@ -542,7 +612,7 @@ async def _aiter_response_stream(self) -> AsyncIterator[str]:
542612
if isinstance(self.response_stream, _HTTPX_RESPONSE_TYPES):
543613
try:
544614
response_stream: Any = self.response_stream
545-
async for line in response_stream.aiter_lines():
615+
async for line in _aiter_sse_lines(response_stream.aiter_bytes()):
546616
if not line:
547617
if data_buffer:
548618
yield '\n'.join(data_buffer)

‎google/genai/tests/client/test_async_stream.py‎

Lines changed: 198 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -16,12 +16,15 @@
1616
"""Tests for async stream."""
1717

1818
import asyncio
19-
from typing import List
19+
import io
20+
import json
21+
from typing import Any, List
2022
from unittest import mock
2123
from unittest.mock import AsyncMock
2224
from unittest.mock import MagicMock
2325
from unittest.mock import patch
2426
import pytest
27+
import requests
2528

2629

2730
try:
@@ -51,15 +54,12 @@ class MockHTTPXResponse(httpx.Response):
5154
"""Mock httpx.Response class for testing."""
5255

5356
def __init__(self, lines: List[str]):
54-
self.aiter_lines = MagicMock()
55-
self.aiter_lines.return_value.__aiter__ = MagicMock(
56-
return_value=self._async_line_iterator(lines)
57-
)
57+
self.aiter_bytes = MagicMock(return_value=self._async_byte_iterator(lines))
5858
self.aclose = AsyncMock()
5959

60-
async def _async_line_iterator(self, lines: List[str]):
60+
async def _async_byte_iterator(self, lines: List[str]):
6161
for line in lines:
62-
yield line
62+
yield line.encode('utf-8') + b'\n'
6363

6464

6565
class MockAIOHTTPResponse(aiohttp.ClientResponse):
@@ -128,7 +128,7 @@ async def test_httpx_simple_lines(responses: api_client.HttpResponse):
128128
results = [line async for line in responses._aiter_response_stream()]
129129

130130
assert results == lines
131-
mock_response.aiter_lines.assert_called_once()
131+
mock_response.aiter_bytes.assert_called_once()
132132
mock_response.aclose.assert_called_once()
133133

134134

@@ -141,7 +141,7 @@ async def test_httpx_data_prefix(responses: api_client.HttpResponse):
141141
results = [line async for line in responses._aiter_response_stream()]
142142

143143
assert results == ["{ 'message': 'hello' }", "{ 'status': 'ok' }"]
144-
mock_response.aiter_lines.assert_called_once()
144+
mock_response.aiter_bytes.assert_called_once()
145145
mock_response.aclose.assert_called_once()
146146

147147
@pytest.mark.asyncio
@@ -158,7 +158,7 @@ async def test_httpx_multiline_data(responses: api_client.HttpResponse) -> None:
158158
results = [line async for line in responses._aiter_response_stream()]
159159

160160
assert results == ["{\n 'message': 'hello'\n}"]
161-
mock_response.aiter_lines.assert_called_once()
161+
mock_response.aiter_bytes.assert_called_once()
162162
mock_response.aclose.assert_called_once()
163163

164164

@@ -170,7 +170,9 @@ def test_httpx_multiline_data_sync(responses: api_client.HttpResponse) -> None:
170170
"",
171171
]
172172
mock_response = MagicMock(spec=httpx.Response)
173-
mock_response.iter_lines.return_value = iter(lines)
173+
mock_response.iter_bytes.return_value = iter(
174+
[b'\n'.join(line.encode('utf-8') for line in lines) + b'\n']
175+
)
174176
responses.response_stream = mock_response
175177

176178
results = list(responses._iter_response_stream())
@@ -193,7 +195,7 @@ async def test_httpx_multiple_json_chunk(responses: api_client.HttpResponse):
193195
results = [line async for line in responses._aiter_response_stream()]
194196

195197
assert results == ['{ "id": 1 }', '{ "id": 2 }', '{ "id": 3 }']
196-
mock_response.aiter_lines.assert_called_once()
198+
mock_response.aiter_bytes.assert_called_once()
197199
mock_response.aclose.assert_called_once()
198200

199201

@@ -207,7 +209,7 @@ async def test_httpx_incomplete_json_at_end(responses: api_client.HttpResponse):
207209

208210
# The remaining chunk is yielded
209211
assert results == ['{ "partial": "data"']
210-
mock_response.aiter_lines.assert_called_once()
212+
mock_response.aiter_bytes.assert_called_once()
211213
mock_response.aclose.assert_called_once()
212214

213215

@@ -220,7 +222,7 @@ async def test_httpx_empty_stream(responses: api_client.HttpResponse):
220222
results = [line async for line in responses._aiter_response_stream()]
221223

222224
assert results == []
223-
mock_response.aiter_lines.assert_called_once()
225+
mock_response.aiter_bytes.assert_called_once()
224226
mock_response.aclose.assert_called_once()
225227

226228

@@ -538,3 +540,185 @@ async def test_error_event_in_generate_content_stream_async(mock_transport):
538540

539541
with pytest.raises(errors.ServerError):
540542
await anext(generated_response)
543+
544+
545+
_LINE_SEPARATORS = pytest.mark.parametrize(
546+
'sep', ['\u2028', '\u2029', '\x85'], ids=['u2028', 'u2029', 'u0085']
547+
)
548+
549+
550+
def _sse_body_with_separator(sep: str) -> bytes:
551+
return b''.join(
552+
b'data: '
553+
+ json.dumps(c, ensure_ascii=False).encode('utf-8')
554+
+ b'\r\n\r\n'
555+
for c in _sse_events_with_separator(sep)
556+
)
557+
558+
559+
def _sse_events_with_separator(sep: str) -> list[dict[str, Any]]:
560+
return [
561+
{
562+
'candidates': [{
563+
'content': {
564+
'role': 'model',
565+
'parts': [{'text': f'line one{sep}line two'}],
566+
}
567+
}]
568+
},
569+
{
570+
'candidates': [{
571+
'content': {'role': 'model', 'parts': [{'text': ' done'}]},
572+
'finishReason': 'STOP',
573+
}]
574+
},
575+
]
576+
577+
578+
def _sse_handler(sep: str):
579+
def handler(_: httpx.Request) -> httpx.Response:
580+
return httpx.Response(
581+
200,
582+
content=_sse_body_with_separator(sep),
583+
headers={'content-type': 'text/event-stream'},
584+
)
585+
586+
return handler
587+
588+
589+
@_LINE_SEPARATORS
590+
def test_generate_content_stream_line_separator_in_text(sep):
591+
"""Unicode line separators inside a JSON string do not split an SSE line."""
592+
client = Client(
593+
api_key='test_api_key',
594+
http_options=types.HttpOptions(
595+
client_args={'transport': httpx.MockTransport(_sse_handler(sep))}
596+
),
597+
)
598+
599+
stream = client.models.generate_content_stream(
600+
model='gemini-2.5-flash', contents='hi'
601+
)
602+
603+
assert [c.text for c in stream] == [f'line one{sep}line two', ' done']
604+
605+
606+
@_LINE_SEPARATORS
607+
@pytest.mark.asyncio
608+
async def test_generate_content_stream_line_separator_in_text_async(sep):
609+
"""Async httpx streams keep Unicode line separators inside JSON strings."""
610+
client = Client(
611+
api_key='test_api_key',
612+
http_options=types.HttpOptions(
613+
async_client_args={
614+
'transport': httpx.MockTransport(_sse_handler(sep))
615+
}
616+
),
617+
)
618+
619+
stream = await client.aio.models.generate_content_stream(
620+
model='gemini-2.5-flash', contents='hi'
621+
)
622+
623+
assert [c.text async for c in stream] == [
624+
f'line one{sep}line two',
625+
' done',
626+
]
627+
628+
629+
@_LINE_SEPARATORS
630+
def test_requests_stream_line_separator_in_text(sep):
631+
"""The requests transport splits SSE lines only at CR, LF and CRLF."""
632+
resp = requests.Response()
633+
resp.status_code = 200
634+
resp.raw = io.BytesIO(_sse_body_with_separator(sep))
635+
resp.encoding = 'utf-8'
636+
637+
results = list(
638+
api_client.HttpResponse(
639+
headers={}, response_stream=resp
640+
)._iter_response_stream()
641+
)
642+
643+
assert results == [
644+
json.dumps(c, ensure_ascii=False) for c in _sse_events_with_separator(sep)
645+
]
646+
647+
648+
def test_httpx_stream_crlf_split_across_chunks():
649+
"""A CRLF split across two chunks ends exactly one line."""
650+
chunks = [b'data: {"a": 1}\r', b'\n\r\n', b'data: {"b": 2}\r\n\r\n']
651+
resp = api_client.HttpResponse(
652+
headers={}, response_stream=mock_response(chunks)
653+
)
654+
655+
assert list(resp._iter_response_stream()) == ['{"a": 1}', '{"b": 2}']
656+
657+
658+
def test_httpx_stream_cr_only_terminators():
659+
"""A bare CR ends a line, as the SSE format requires."""
660+
chunks = [b'data: {"a": 1}\r\rdata: {"b": 2}\r\r']
661+
resp = api_client.HttpResponse(
662+
headers={}, response_stream=mock_response(chunks)
663+
)
664+
665+
assert list(resp._iter_response_stream()) == ['{"a": 1}', '{"b": 2}']
666+
667+
668+
def test_httpx_stream_multibyte_char_split_across_chunks():
669+
"""A UTF-8 character split across chunks decodes as one character."""
670+
chunks = [b'data: {"t": "x\xe2\x80', b'\xa8y"}\n\n']
671+
resp = api_client.HttpResponse(
672+
headers={}, response_stream=mock_response(chunks)
673+
)
674+
675+
assert list(resp._iter_response_stream()) == ['{"t": "x\u2028y"}']
676+
677+
678+
def test_httpx_stream_multiline_data_with_line_separator():
679+
"""Multi-line data fields still join with LF around a U+2028."""
680+
chunks = [b'data: {"t": "a\xe2\x80\xa8b",\n', b'data: "u": 1}\n\n']
681+
resp = api_client.HttpResponse(
682+
headers={}, response_stream=mock_response(chunks)
683+
)
684+
685+
assert list(resp._iter_response_stream()) == ['{"t": "a\u2028b",\n"u": 1}']
686+
687+
688+
def test_httpx_stream_invalid_utf8_is_replaced():
689+
"""Invalid UTF-8 bytes become U+FFFD instead of raising."""
690+
chunks = [b'data: {"t": "\xff"}\n\n']
691+
resp = api_client.HttpResponse(
692+
headers={}, response_stream=mock_response(chunks)
693+
)
694+
695+
assert list(resp._iter_response_stream()) == ['{"t": "\ufffd"}']
696+
697+
698+
def test_httpx_stream_final_line_without_terminator():
699+
"""A last line with no terminator is still yielded at end of stream."""
700+
chunks = [b'data: {"a": 1}\n\n{"error": ', b'{"code": 500}}']
701+
resp = api_client.HttpResponse(
702+
headers={}, response_stream=mock_response(chunks)
703+
)
704+
705+
assert list(resp._iter_response_stream()) == [
706+
'{"a": 1}',
707+
'{"error": {"code": 500}}',
708+
]
709+
710+
711+
@pytest.mark.asyncio
712+
async def test_httpx_async_stream_line_separator_in_text(
713+
responses: api_client.HttpResponse,
714+
):
715+
"""The async httpx branch keeps a U+2028 line whole and closes once."""
716+
httpx_response = MockHTTPXResponse(
717+
['data: {"t": "a\u2028b"}', '', 'data: {"t": "c"}', '']
718+
)
719+
responses.response_stream = httpx_response
720+
721+
results = [line async for line in responses._aiter_response_stream()]
722+
723+
assert results == ['{"t": "a\u2028b"}', '{"t": "c"}']
724+
httpx_response.aclose.assert_called_once()

0 commit comments

Comments
 (0)