diff --git a/google/genai/_api_client.py b/google/genai/_api_client.py index 36d7a2638d..c229e5a6cb 100644 --- a/google/genai/_api_client.py +++ b/google/genai/_api_client.py @@ -362,6 +362,23 @@ class HttpRequest: timeout: Optional[float] = None +async def _aiter_aiohttp_lines( + content: 'aiohttp.StreamReader', +) -> AsyncIterator[bytes]: + """Yields newline-terminated lines from an aiohttp stream, of any length.""" + buffer = bytearray() + async for data in content.iter_any(): + scan_start = len(buffer) + buffer += data + line_start = 0 + while (newline_pos := buffer.find(b'\n', scan_start)) != -1: + yield bytes(buffer[line_start : newline_pos + 1]) + line_start = scan_start = newline_pos + 1 + del buffer[:line_start] + if buffer: + yield bytes(buffer) + + class HttpResponse: def __init__( @@ -581,18 +598,9 @@ async def _aiter_response_stream(self) -> AsyncIterator[str]: self.response_stream, aiohttp.ClientResponse ): try: - while True: - # Read a line from the stream. This returns bytes. - try: - line_bytes = await self.response_stream.content.readline( - max_line_length=READ_BUFFER_SIZE # type: ignore[call-arg] - ) - except TypeError: - # Ensure backwards compatibility with older versions of - # aiohttp that do not support max_line_length. - line_bytes = await self.response_stream.content.readline() - if not line_bytes: - break + async for line_bytes in _aiter_aiohttp_lines( + self.response_stream.content + ): # Decode the bytes and remove trailing whitespace and newlines. line = line_bytes.decode('utf-8').rstrip() if not line: diff --git a/google/genai/tests/client/test_async_stream.py b/google/genai/tests/client/test_async_stream.py index 71b626b9f7..2fce77cdd2 100644 --- a/google/genai/tests/client/test_async_stream.py +++ b/google/genai/tests/client/test_async_stream.py @@ -62,30 +62,36 @@ async def _async_line_iterator(self, lines: List[str]): yield line +def _aiohttp_stream_reader( + data: bytes, piece_size: int = 65536, eof: bool = True +): + """Returns a real aiohttp StreamReader fed with data in pieces.""" + reader = aiohttp.StreamReader( + mock.Mock(_reading_paused=False), + 2**16, + loop=asyncio.get_running_loop(), + ) + for i in range(0, len(data), piece_size): + reader.feed_data(data[i : i + piece_size]) + if eof: + reader.feed_eof() + return reader + + class MockAIOHTTPResponse(aiohttp.ClientResponse): def __init__(self, lines: List[str]): - self.content = MagicMock() - self.content.readline = AsyncMock() - # Simulate reading lines, each ending with newline bytes for readline behavior - self._read_data = b"\n".join(line.encode("utf-8") for line in lines) + b"\n" - self._read_pos = 0 - self.content.readline.side_effect = self._async_read_line + self.content = _aiohttp_stream_reader( + b'\n'.join(line.encode('utf-8') for line in lines) + b'\n' + ) self.release = MagicMock() - async def _async_read_line(self, **kwargs) -> bytes: - if self._read_pos >= len(self._read_data): - return b"" # End of stream - newline_pos = self._read_data.find(b"\n", self._read_pos) - if newline_pos == -1: # Should not happen with the appended '\n' - line = self._read_data[self._read_pos :] - self._read_pos = len(self._read_data) - return line - else: - line = self._read_data[self._read_pos : newline_pos + 1] - self._read_pos = newline_pos + 1 - return line +class MockAIOHTTPResponseFromReader(aiohttp.ClientResponse): + + def __init__(self, content): + self.content = content + self.release = MagicMock() @pytest.fixture @@ -237,7 +243,6 @@ async def test_aiohttp_simple_lines(responses: api_client.HttpResponse): results = [line async for line in responses._aiter_response_stream()] assert results == lines - mock_response.content.readline.assert_called() mock_response.release.assert_called_once() @@ -253,7 +258,6 @@ async def test_aiohttp_data_prefix(responses: api_client.HttpResponse): results = [line async for line in responses._aiter_response_stream()] assert results == ["{ 'message': 'hello' }", "{ 'status': 'ok' }"] - mock_response.content.readline.assert_called() mock_response.release.assert_called_once() @@ -275,7 +279,6 @@ async def test_aiohttp_multiple_json_chunks(responses: api_client.HttpResponse): results = [line async for line in responses._aiter_response_stream()] assert results == ['{ "id": 1 }', '{ "id": 2 }', '{ "id": 3 }'] - mock_response.content.readline.assert_called() mock_response.release.assert_called_once() @@ -293,51 +296,9 @@ async def test_aiohttp_incomplete_json_at_end( results = [line async for line in responses._aiter_response_stream()] assert results == ['{ "partial": "data"'] - mock_response.content.readline.assert_called() mock_response.release.assert_called_once() -class MockAIOHTTPResponseWithLineLimits(aiohttp.ClientResponse): - """Mock that enforces aiohttp's real readline limits. - - Real aiohttp StreamReader raises LineTooLong when a line exceeds - `_high_water` (= limit * 2) bytes. The default limit is 2**16, so lines - over 131072 bytes fail unless `max_line_length` is explicitly passed. - """ - - DEFAULT_HIGH_WATER = 2**16 * 2 # 131072, same as aiohttp default - - def __init__(self, lines: List[str]): - self.content = MagicMock() - self.content.readline = AsyncMock() - self._read_data = b"\n".join(line.encode("utf-8") for line in lines) + b"\n" - self._read_pos = 0 - self.content.readline.side_effect = self._async_read_line - self.release = MagicMock() - - async def _async_read_line( - self, *, max_line_length=None - ) -> bytes: - if self._read_pos >= len(self._read_data): - return b"" - - newline_pos = self._read_data.find(b"\n", self._read_pos) - if newline_pos == -1: - line = self._read_data[self._read_pos:] - self._read_pos = len(self._read_data) - else: - line = self._read_data[self._read_pos:newline_pos + 1] - self._read_pos = newline_pos + 1 - - # Enforce limit like real aiohttp StreamReader.readuntil does - limit = max_line_length or self.DEFAULT_HIGH_WATER - if len(line) > limit: - from aiohttp.http_exceptions import LineTooLong - raise LineTooLong(line[:100] + b"...", limit) - - return line - - @requires_aiohttp @pytest.mark.asyncio async def test_aiohttp_large_sse_line_with_thought_signature( @@ -350,12 +311,8 @@ async def test_aiohttp_large_sse_line_with_thought_signature( of 131072 bytes. Thinking models can return a thoughtSignature field large enough to push a single SSE data: line past this limit, causing LineTooLong. - The fix passes max_line_length=READ_BUFFER_SIZE (4MB) explicitly on the - readline() call in _aiter_response_stream(), overriding the limit at the - call site regardless of how the underlying aiohttp session was configured. - - This test verifies the fix by using a mock that enforces the real aiohttp - readline limit and confirms a 150KB line is streamed successfully. + This test uses a real aiohttp StreamReader and confirms a 150KB line is + streamed successfully. """ api_client.has_aiohttp = True @@ -367,7 +324,7 @@ async def test_aiohttp_large_sse_line_with_thought_signature( ) lines = [f"data: {large_sse_payload}", ""] - mock_response = MockAIOHTTPResponseWithLineLimits(lines) + mock_response = MockAIOHTTPResponse(lines) responses.response_stream = mock_response results = [line async for line in responses._aiter_response_stream()] @@ -378,6 +335,106 @@ async def test_aiohttp_large_sse_line_with_thought_signature( mock_response.release.assert_called_once() +@requires_aiohttp +@pytest.mark.asyncio +async def test_aiohttp_sse_line_larger_than_read_buffer_size( + responses: api_client.HttpResponse, +): + """Verifies an SSE line over every aiohttp line cap is delivered whole. + + The payload exceeds both READ_BUFFER_SIZE and the 2 * read_bufsize high-water + mark of the session the SDK creates. + """ + api_client.has_aiohttp = True + payload = ( + '{"inlineData": {"data": "' + + 'A' * (2 * api_client.READ_BUFFER_SIZE + 1) + + '"}}' + ) + body = ( + b'data: {"text": "hi"}\r\n\r\n' + + b'data: ' + + payload.encode() + + b'\r\n\r\n' + ) + mock_response = MockAIOHTTPResponseFromReader(_aiohttp_stream_reader(body)) + responses.response_stream = mock_response + + results = [line async for line in responses._aiter_response_stream()] + + assert results == ['{"text": "hi"}', payload] + mock_response.release.assert_called_once() + + +@requires_aiohttp +@pytest.mark.asyncio +async def test_aiohttp_lines_split_across_reads( + responses: api_client.HttpResponse, +): + """Verifies lines are reassembled when every byte arrives in its own read.""" + api_client.has_aiohttp = True + body = b'data: { "id": 1 }\r\n\r\ndata: { "id": 2 }\r\n\r\n' + mock_response = MockAIOHTTPResponseFromReader( + _aiohttp_stream_reader(body, piece_size=1) + ) + responses.response_stream = mock_response + + results = [line async for line in responses._aiter_response_stream()] + + assert results == ['{ "id": 1 }', '{ "id": 2 }'] + mock_response.release.assert_called_once() + + +@requires_aiohttp +@pytest.mark.asyncio +async def test_aiohttp_trailing_line_without_newline( + responses: api_client.HttpResponse, +): + """Verifies a final line with no newline is still yielded.""" + api_client.has_aiohttp = True + mock_response = MockAIOHTTPResponseFromReader( + _aiohttp_stream_reader(b'{ "id": 1 }') + ) + responses.response_stream = mock_response + + results = [line async for line in responses._aiter_response_stream()] + + assert results == ['{ "id": 1 }'] + mock_response.release.assert_called_once() + + +@requires_aiohttp +@pytest.mark.asyncio +async def test_aiohttp_empty_stream(responses: api_client.HttpResponse): + """Verifies an empty body yields nothing and releases the response.""" + api_client.has_aiohttp = True + mock_response = MockAIOHTTPResponseFromReader(_aiohttp_stream_reader(b'')) + responses.response_stream = mock_response + + results = [line async for line in responses._aiter_response_stream()] + + assert results == [] + mock_response.release.assert_called_once() + + +@requires_aiohttp +@pytest.mark.asyncio +async def test_aiohttp_stream_error_releases_response( + responses: api_client.HttpResponse, +): + """Verifies a transport error propagates and the response is released.""" + api_client.has_aiohttp = True + reader = _aiohttp_stream_reader(b'data: { "id": 1 }\r\n', eof=False) + reader.set_exception(aiohttp.ClientPayloadError('boom')) + mock_response = MockAIOHTTPResponseFromReader(reader) + responses.response_stream = mock_response + + with pytest.raises(aiohttp.ClientPayloadError, match='boom'): + async for _ in responses._aiter_response_stream(): + pass + mock_response.release.assert_called_once() + + def mock_response(chunks): mock_stream = MagicMock(spec=httpx.SyncByteStream) mock_stream.__iter__.return_value = chunks diff --git a/google/genai/tests/client/test_retries.py b/google/genai/tests/client/test_retries.py index 6332f2091f..1198da0708 100644 --- a/google/genai/tests/client/test_retries.py +++ b/google/genai/tests/client/test_retries.py @@ -662,8 +662,8 @@ async def _aiohttp_async_response(status: int, streamable: bool = False): response.json.return_value = {} response.text.return_value = 'test' if streamable: - response.content = mock.Mock() - response.content.readline = mock.AsyncMock(return_value=b'') + response.content = mock.MagicMock() + response.content.iter_any.return_value.__aiter__.return_value = [] response.release = mock.MagicMock() return response