From a207ae7ce1d29e9c07376d98f25c2490c3558fdd Mon Sep 17 00:00:00 2001 From: Rob Clevenger Date: Wed, 11 Nov 2026 20:06:19 +0000 Subject: [PATCH] fix: close async stream when an error event arrives mid-stream async_request_streamed raises APIError from inside its `async for` over the response when the server sends an `{"error": ...}` event. Leaving an `async for` by exception does not close the iterator, so HttpResponse.async_segments and the _aiter_response_stream it reads stay suspended. The connection is released only when the garbage collector finalizes them, which asyncio does later in a task of its own, outside the Context of the request that failed. Wrap both loops in contextlib.aclosing, so the stream is closed before the error reaches the caller. Annotate the two generators as AsyncGenerator so the type checker accepts them as closeable. --- google/genai/_api_client.py | 49 ++++++++++-------- .../genai/tests/client/test_async_stream.py | 51 +++++++++++++++++++ 2 files changed, 80 insertions(+), 20 deletions(-) diff --git a/google/genai/_api_client.py b/google/genai/_api_client.py index d8a49a5736..da3abb73a0 100644 --- a/google/genai/_api_client.py +++ b/google/genai/_api_client.py @@ -20,7 +20,9 @@ """ import asyncio +from collections.abc import AsyncGenerator from collections.abc import Generator +import contextlib import copy from dataclasses import dataclass import inspect @@ -462,7 +464,7 @@ def segments(self) -> Generator[Any, None, None]: for chunk in self._iter_response_stream(): yield self._load_json_from_response(chunk) - async def async_segments(self) -> AsyncIterator[Any]: + async def async_segments(self) -> AsyncGenerator[Any, None]: if isinstance(self.response_stream, list): # list of objects retrieved from replay or from non-streaming API. for chunk in self.response_stream: @@ -471,9 +473,12 @@ async def async_segments(self) -> AsyncIterator[Any]: async for c in []: # type: ignore[attr-defined] yield c else: - # Iterator of objects retrieved from the API. - async for chunk in self._aiter_response_stream(): - yield self._load_json_from_response(chunk) + # Iterator of objects retrieved from the API. Close the stream if this + # generator is closed early, so the connection is released here rather + # than when the garbage collector finalizes a suspended generator. + async with contextlib.aclosing(self._aiter_response_stream()) as stream: + async for chunk in stream: + yield self._load_json_from_response(chunk) def byte_segments(self) -> Generator[Union[bytes, Any], None, None]: if isinstance(self.byte_stream, list): @@ -548,7 +553,7 @@ def _iter_response_stream(self) -> Iterator[str]: if data_buffer: yield '\n'.join(data_buffer) - async def _aiter_response_stream(self) -> AsyncIterator[str]: + async def _aiter_response_stream(self) -> AsyncGenerator[str, None]: """Asynchronously iterates over chunks retrieved from the API.""" is_valid_response = isinstance( self.response_stream, _HTTPX_RESPONSE_TYPES @@ -2244,22 +2249,26 @@ async def async_request_streamed( ) async def async_generator(): # type: ignore[no-untyped-def] - async for chunk in response: - chunk_dump = json.dumps(chunk) - try: - if chunk_dump.startswith('{"error":'): - chunk_json = json.loads(chunk_dump) - await errors.APIError.raise_error_async( - chunk_json.get('error', {}).get('code'), - chunk_json, - response, + # An error event raises from inside this loop. Close the segments on the + # way out, so the stream is released before the error reaches the caller. + async with contextlib.aclosing(response.async_segments()) as segments: + async for chunk in segments: + chunk_dump = json.dumps(chunk) + try: + if chunk_dump.startswith('{"error":'): + chunk_json = json.loads(chunk_dump) + await errors.APIError.raise_error_async( + chunk_json.get('error', {}).get('code'), + chunk_json, + response, + ) + except json.decoder.JSONDecodeError: + logger.debug( + 'Failed to decode chunk that contains an error: %s' + % chunk_dump ) - except json.decoder.JSONDecodeError: - logger.debug( - 'Failed to decode chunk that contains an error: %s' % chunk_dump - ) - pass - yield SdkHttpResponse(headers=response.headers, body=chunk_dump) + pass + yield SdkHttpResponse(headers=response.headers, body=chunk_dump) return async_generator() # type: ignore[no-untyped-call] diff --git a/google/genai/tests/client/test_async_stream.py b/google/genai/tests/client/test_async_stream.py index 43b57eb033..1f633a9405 100644 --- a/google/genai/tests/client/test_async_stream.py +++ b/google/genai/tests/client/test_async_stream.py @@ -545,3 +545,54 @@ async def test_error_event_in_generate_content_stream_async(mock_transport): with pytest.raises(errors.ServerError): await anext(generated_response) + + +@patch.object(httpx, 'AsyncBaseTransport') +@pytest.mark.asyncio +async def test_error_event_in_streamed_responses_async_closes_stream( + mock_transport, +): + """An error event mid-stream closes the HTTP stream before the error surfaces. + + The server sends the error event and leaves the connection open. The SDK + raises from inside the loop that reads the stream, so unless that loop closes + what it reads, the stream generators stay suspended until garbage collection, + and asyncio finalizes them later in a task of its own. The stream must already + be closed when the caller sees the ServerError. + """ + mock_stream = MagicMock(spec=httpx.AsyncByteStream) + + async def _aiter(): + yield EVENT_STREAM_DATA_WITH_ERROR[0] + yield EVENT_STREAM_DATA_WITH_ERROR[1] + yield EVENT_STREAM_DATA_WITH_ERROR[2] + b'\n' + # The server has not ended the stream, so nothing closes it on exhaustion. + await asyncio.Event().wait() + + mock_stream.__aiter__.side_effect = _aiter + mock_stream.aclose = AsyncMock() + + async def _handle_async_request(_): + return httpx.Response(status_code=200, stream=mock_stream) + + client = api_client.BaseApiClient( + api_key='test_api_key', + http_options=types.HttpOptions( + async_client_args={'transport': mock_transport} + ), + ) + mock_transport.handle_async_request = _handle_async_request + mock_transport.aclose = AsyncMock() + + resp = await client.async_request_streamed( + 'POST', 'models/gemini-2.5-flash', {'key': 'value'} + ) + await anext(resp) + mock_stream.aclose.assert_not_awaited() + + with pytest.raises(errors.ServerError): + await anext(resp) + + # No await between the raise and this check, so a close scheduled by the + # event loop's async generator finalizer cannot have run yet. + mock_stream.aclose.assert_awaited_once()