fix: close async stream when an error event arrives mid-stream

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]
 

google/genai/tests/client/test_async_stream.py

@@ -539,3 +539,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()