fix(streaming): close the SSE event iterators a stream opens

src/openai/_streaming.py

@@ -13,7 +13,7 @@
 from ._utils import is_mapping, extract_type_var_from_base
 from ._httpx2 import request_exceptions, timeout_exceptions
 from ._exceptions import APIError, APITimeoutError, APIConnectionError
-from ._streaming_cleanup import drain_response_async
+from ._streaming_cleanup import close_owned, run_cleanup, aclose_owned, drain_response_async
 
 if TYPE_CHECKING:
     from ._client import OpenAI, AsyncOpenAI
@@ -49,22 +49,39 @@ def __next__(self) -> _T:
         return self._iterator.__next__()
 
     def __iter__(self) -> Iterator[_T]:
-        for item in self._iterator:
-            yield item
+        return self._iterator
 
     def _iter_events(self) -> Iterator[ServerSentEvent]:
+        chunks: Optional[Iterator[bytes]] = None
+        events: Optional[Iterator[ServerSentEvent]] = None
+        failure: Optional[BaseException] = None
         try:
-            yield from self._decoder.iter_bytes(self.response.iter_bytes())
-        except timeout_exceptions() as err:
-            raise APITimeoutError(request=self.response.request) from err
-        except request_exceptions() as err:
-            raise APIConnectionError(request=self.response.request) from err
+            try:
+                chunks = self.response.iter_bytes()
+                events = self._decoder.iter_bytes(chunks)
+                # Not `yield from`: it would close `events` itself, and the
+                # cleanup below must stay the only close.
+                for sse in events:
+                    yield sse
+            except timeout_exceptions() as err:
+                raise APITimeoutError(request=self.response.request) from err
+            except request_exceptions() as err:
+                raise APIConnectionError(request=self.response.request) from err
+        except GeneratorExit:
+            raise  # Normal closure, not a failure that should outrank cleanup errors.
+        except BaseException as err:
+            failure = err
+            raise
+        finally:
+            # Close the iterators opened here; stopping early does not close them.
+            close_owned(events, chunks, primary=failure)
 
     def __stream__(self) -> Iterator[_T]:
         cast_to = cast(Any, self._cast_to)
         response = self.response
         process_data = self._client._process_response_data
         iterator = self._iter_events()
+        failure: Optional[BaseException] = None
 
         try:
             for sse in iterator:
@@ -113,9 +130,15 @@ def __stream__(self) -> Iterator[_T]:
                         cast_to=cast_to,
                         response=response,
                     )
+        except GeneratorExit:
+            raise  # Normal closure, not a failure that should outrank cleanup errors.
+        except BaseException as err:
+            failure = err
+            raise
         finally:
-            # Ensure the response is closed even if the consumer doesn't read all data
-            response.close()
+            # Close the event iterator explicitly; stopping at [DONE] leaves it suspended.
+            # Ensure the response is closed even if the consumer doesn't read all data.
+            close_owned(iterator, response, primary=failure)
 
     def __enter__(self) -> Self:
         return self
@@ -126,15 +149,19 @@ def __exit__(
         exc: BaseException | None,
         exc_tb: TracebackType | None,
     ) -> None:
-        self.close()
+        # A cleanup error must not replace the exception leaving the block.
+        close_owned(self, primary=exc)
 
     def close(self) -> None:
         """
         Close the response and release the connection.
 
         Automatically called if the response body is read to completion.
         """
-        self.response.close()
+        # Concurrent close does not cancel or join an active consumer; that consumer
+        # closes the iterator when it unwinds.
+        running = getattr(self._iterator, "gi_running", False)
+        close_owned(None if running else self._iterator, self.response)
 
 
 class AsyncStream(Generic[_T]):
@@ -162,24 +189,38 @@ def __init__(
     async def __anext__(self) -> _T:
         return await self._iterator.__anext__()
 
-    async def __aiter__(self) -> AsyncIterator[_T]:
-        async for item in self._iterator:
-            yield item
+    def __aiter__(self) -> AsyncIterator[_T]:
+        return self._iterator
 
     async def _iter_events(self) -> AsyncIterator[ServerSentEvent]:
+        chunks: Optional[AsyncIterator[bytes]] = None
+        events: Optional[AsyncIterator[ServerSentEvent]] = None
+        failure: Optional[BaseException] = None
         try:
-            async for sse in self._decoder.aiter_bytes(self.response.aiter_bytes()):
-                yield sse
-        except timeout_exceptions() as err:
-            raise APITimeoutError(request=self.response.request) from err
-        except request_exceptions() as err:
-            raise APIConnectionError(request=self.response.request) from err
+            try:
+                chunks = self.response.aiter_bytes()
+                events = self._decoder.aiter_bytes(chunks)
+                async for sse in events:
+                    yield sse
+            except timeout_exceptions() as err:
+                raise APITimeoutError(request=self.response.request) from err
+            except request_exceptions() as err:
+                raise APIConnectionError(request=self.response.request) from err
+        except GeneratorExit:
+            raise  # Normal closure, not a failure that should outrank cleanup errors.
+        except BaseException as err:
+            failure = err
+            raise
+        finally:
+            # Close the iterators opened here; an interrupted `async for` does not.
+            await aclose_owned(events, chunks, primary=failure)
 
     async def __stream__(self) -> AsyncIterator[_T]:
         cast_to = cast(Any, self._cast_to)
         response = self.response
         process_data = self._client._process_response_data
         iterator = self._iter_events()
+        failure: Optional[BaseException] = None
 
         try:
             async for sse in iterator:
@@ -233,9 +274,17 @@ async def __stream__(self) -> AsyncIterator[_T]:
                         cast_to=cast_to,
                         response=response,
                     )
+        except GeneratorExit:
+            raise  # Normal closure, not a failure that should outrank cleanup errors.
+        except BaseException as err:
+            failure = err
+            raise
         finally:
-            # Ensure the response is closed even if the consumer doesn't read all data
-            await response.aclose()
+            # Close the event iterator explicitly. Stopping at [DONE] leaves it and the
+            # httpx2 byte iterators beneath it suspended; left to the garbage collector,
+            # asyncio closes those nested generators concurrently.
+            # Ensure the response is closed even if the consumer doesn't read all data.
+            await aclose_owned(iterator, response, primary=failure)
 
     async def __aenter__(self) -> Self:
         return self
@@ -246,15 +295,19 @@ async def __aexit__(
         exc: BaseException | None,
         exc_tb: TracebackType | None,
     ) -> None:
-        await self.close()
+        # A cleanup error must not replace the exception leaving the block.
+        await run_cleanup(self.close, primary=exc)
 
     async def close(self) -> None:
         """
         Close the response and release the connection.
 
         Automatically called if the response body is read to completion.
         """
-        await self.response.aclose()
+        # Concurrent close does not cancel or join an active consumer; that consumer
+        # closes the iterator when it unwinds.
+        running = getattr(self._iterator, "ag_running", False)
+        await aclose_owned(None if running else self._iterator, self.response)
 
     async def aclose(self) -> None:
         """Close the response and release the connection. Alias for `close()`."""

src/openai/_streaming_cleanup.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from typing import AsyncIterator, cast
+from typing import Any, Callable, Optional, Awaitable, AsyncIterator, cast
 
 import anyio
 import httpx2
@@ -27,3 +27,83 @@ async def drain_response_async(response: httpx2.Response) -> None:
         # cancellation. Cancellation still propagates after the response closes.
         with anyio.CancelScope(shield=True):
             await response.aclose()
+
+
+def _cleanup_error(primary: Optional[BaseException], failure: Optional[BaseException]) -> Optional[BaseException]:
+    """Pick the cleanup failure to raise, if any.
+
+    A failure already propagating wins, except that a control-flow exception raised
+    during cleanup (for example cancellation) is never dropped in favor of an ordinary
+    error.
+    """
+    if failure is None:
+        return None
+    if primary is None:
+        return failure
+    if not isinstance(failure, Exception) and isinstance(primary, Exception):
+        return failure
+    return None
+
+
+def _keep(failure: Optional[BaseException], err: BaseException) -> BaseException:
+    """Keep the first cleanup failure, but let a later control-flow exception replace an ordinary one."""
+    if failure is None or (isinstance(failure, Exception) and not isinstance(err, Exception)):
+        return err
+    return failure
+
+
+def close_owned(*resources: object, primary: Optional[BaseException] = None) -> None:
+    """Close resources this caller owns, in order.
+
+    Every resource gets its close attempt; objects without ``close()`` are skipped.
+    The first cleanup failure is raised after all attempts unless ``primary``, the
+    failure already propagating, takes precedence (see ``_cleanup_error``).
+    """
+    failure: Optional[BaseException] = None
+    for resource in resources:
+        try:
+            close = getattr(resource, "close", None)
+            if close is not None:
+                close()
+        except BaseException as err:  # re-raised below, after the remaining closes
+            failure = _keep(failure, err)
+    error = _cleanup_error(primary, failure)
+    if error is not None:
+        raise error
+
+
+async def aclose_owned(*resources: object, primary: Optional[BaseException] = None) -> None:
+    """Async counterpart of ``close_owned``.
+
+    The closes are shielded from AnyIO cancellation so an already-cancelled scope cannot
+    leave a resource half closed; cancellation still propagates afterwards. The shield
+    has no deadline, so a closer that never returns blocks the caller.
+    """
+    failure: Optional[BaseException] = None
+    with anyio.CancelScope(shield=True):
+        for resource in resources:
+            try:
+                aclose = getattr(resource, "aclose", None)
+                if aclose is not None:
+                    await aclose()
+            except BaseException as err:  # re-raised below, after the remaining closes
+                failure = _keep(failure, err)
+    error = _cleanup_error(primary, failure)
+    if error is not None:
+        raise error
+
+
+async def run_cleanup(cleanup: Callable[[], Awaitable[Any]], *, primary: Optional[BaseException] = None) -> None:
+    """Run one async cleanup callable with the shield and error precedence of ``aclose_owned``.
+
+    For owners whose async close method is not named ``aclose``.
+    """
+    failure: Optional[BaseException] = None
+    with anyio.CancelScope(shield=True):
+        try:
+            await cleanup()
+        except BaseException as err:  # re-raised below according to precedence
+            failure = err
+    error = _cleanup_error(primary, failure)
+    if error is not None:
+        raise error