fix(streaming): ensure remaining body is consumed after [DONE] in Str…

src/openai/_streaming.py

@@ -10,7 +10,7 @@
 
 import httpx2
 
-from ._utils import is_mapping, extract_type_var_from_base
+from ._utils import is_mapping, drain_async_iterator, extract_type_var_from_base
 from ._exceptions import APIError
 
 if TYPE_CHECKING:
@@ -150,6 +150,8 @@ def __init__(
         self._client = client
         self._options = options
         self._decoder = client._make_sse_decoder()
+        self._byte_iterator: AsyncIterator[bytes] | None = None
+        self._done_seen = False
         self._iterator = self.__stream__()
 
     async def __anext__(self) -> _T:
@@ -160,18 +162,22 @@ async def __aiter__(self) -> AsyncIterator[_T]:
             yield item
 
     async def _iter_events(self) -> AsyncIterator[ServerSentEvent]:
-        async for sse in self._decoder.aiter_bytes(self.response.aiter_bytes()):
+        if self._byte_iterator is None:
+            self._byte_iterator = self.response.aiter_bytes()
+        async for sse in self._decoder.aiter_bytes(self._byte_iterator):
             yield sse
 
     async def __stream__(self) -> AsyncIterator[_T]:
         cast_to = cast(Any, self._cast_to)
         response = self.response
         process_data = self._client._process_response_data
+        self._byte_iterator = response.aiter_bytes()
         iterator = self._iter_events()
 
         try:
             async for sse in iterator:
                 if sse.data.startswith("[DONE]"):
+                    self._done_seen = True
                     break
 
                 # we have to special case the Assistants `thread.` events since we won't have an "event" key in the data
@@ -217,8 +223,13 @@ async def __stream__(self) -> AsyncIterator[_T]:
                         response=response,
                     )
         finally:
-            # Ensure the response is closed even if the consumer doesn't read all data
-            await response.aclose()
+            # Bounds cancellation-cooperative async drain for connection reuse.
+            if self._done_seen:
+                await drain_async_iterator(self._byte_iterator, response=response, timeout_ms=50)
+            try:
+                await response.aclose()
+            except Exception:
+                pass
 
     async def __aenter__(self) -> Self:
         return self

src/openai/_utils/__init__.py

@@ -52,7 +52,11 @@
     strip_annotated_type as strip_annotated_type,
     extract_type_var_from_base as extract_type_var_from_base,
 )
-from ._streams import consume_sync_iterator as consume_sync_iterator, consume_async_iterator as consume_async_iterator
+from ._streams import (
+    drain_async_iterator as drain_async_iterator,
+    consume_sync_iterator as consume_sync_iterator,
+    consume_async_iterator as consume_async_iterator,
+)
 from ._transform import (
     PropertyInfo as PropertyInfo,
     transform as transform,

src/openai/_utils/_streams.py

@@ -1,6 +1,8 @@
-from typing import Any
+from typing import Any, Optional
 from typing_extensions import Iterator, AsyncIterator
 
+import anyio
+
 
 def consume_sync_iterator(iterator: Iterator[Any]) -> None:
     for _ in iterator:
@@ -10,3 +12,31 @@ def consume_sync_iterator(iterator: Iterator[Any]) -> None:
 async def consume_async_iterator(iterator: AsyncIterator[Any]) -> None:
     async for _ in iterator:
         ...
+
+
+async def drain_async_iterator(
+    iterator: Optional[AsyncIterator[Any]], response: Any = None, timeout_ms: int = 50
+) -> None:
+    """Drain trailing bytes from async iterator with bounded timeout.
+
+    Bounds cancellation-cooperative async streams to enable connection reuse after [DONE].
+    Uses anyio for backend compatibility (asyncio and Trio). Cannot guarantee termination
+    for custom iterators that block the event loop or suppress cancellation.
+    """
+    if iterator is None:
+        return
+
+    try:
+        with anyio.move_on_after(timeout_ms / 1000.0):
+            try:
+                async for _ in iterator:
+                    pass
+            except Exception:
+                pass
+    finally:
+        # Close response to interrupt any blocked iterator reads on timeout
+        if response is not None:
+            try:
+                await response.aclose()
+            except Exception:
+                pass

tests/test_streaming.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
-from typing import Iterator, AsyncIterator
+from typing import Any
+from collections.abc import Iterator, AsyncIterator
 
 import httpx2
 import pytest
@@ -216,6 +217,214 @@ def body() -> Iterator[bytes]:
     assert sse.json() == {"content": "известни"}
 
 
+def test_done_closes_response_sync(client: OpenAI) -> None:
+    """Sync stream closes response after [DONE]."""
+
+    def body() -> Iterator[bytes]:
+        yield b'data: {"foo":true}\n\n'
+        yield b"data: [DONE]\n\n"
+        yield b": trailing comment after done\n\n"
+
+    response = httpx2.Response(200, content=body())
+    stream: Stream[object] | AsyncStream[object] = Stream(cast_to=object, client=client, response=response)
+    chunks = list(stream)
+
+    assert chunks == [{"foo": True}]
+    assert response.is_closed is True
+
+
+@pytest.mark.asyncio
+async def test_done_drains_remaining_body_async(async_client: AsyncOpenAI) -> None:
+    """After [DONE], async drains remaining body for connection reuse."""
+    exhausted = False
+
+    def body() -> Iterator[bytes]:
+        nonlocal exhausted
+        yield b'data: {"foo":true}\n\n'
+        yield b"data: [DONE]\n\n"
+        yield b": trailing comment after done\n\n"
+        exhausted = True
+
+    async_body = to_aiter(body())
+    response = httpx2.Response(200, content=async_body)
+    stream = AsyncStream(cast_to=object, client=async_client, response=response)
+    chunks = [chunk async for chunk in stream]
+
+    assert chunks == [{"foo": True}]
+    assert exhausted is True
+    assert response.is_closed is True
+
+
+@pytest.mark.asyncio
+async def test_early_exit_without_done_doesnt_drain_async(async_client: AsyncOpenAI) -> None:
+    """Early exit before [DONE] should not consume trailing body and close promptly."""
+    drained = False
+    exhausted = False
+
+    async def patched_drain(*_args: Any, **_kwargs: Any) -> None:
+        """Track if drain was called."""
+        nonlocal drained
+        drained = True
+
+    def body() -> Iterator[bytes]:
+        nonlocal exhausted
+        yield b'data: {"foo":true}\n\n'
+        # No [DONE] sent, consumer will exit early
+        yield b": trailing comment that should not be consumed\n\n"
+        exhausted = True
+
+    async_body = to_aiter(body())
+    response = httpx2.Response(200, content=async_body)
+
+    # Patch drain_async_iterator to track if it's called
+    import openai._streaming as streaming_module
+
+    original_drain = streaming_module.drain_async_iterator
+    streaming_module.drain_async_iterator = patched_drain
+
+    try:
+        stream = AsyncStream(cast_to=object, client=async_client, response=response)
+
+        # Only consume the first chunk, exit early before [DONE]
+        first_chunk = await stream.__anext__()
+        assert first_chunk == {"foo": True}
+
+        # Trailing body should NOT be exhausted since _done_seen is False
+        assert exhausted is False
+
+        # Explicitly close the generator to trigger finally block
+        await stream._iterator.aclose()  # type: ignore[attr-defined]
+
+        # Drain should NOT have been called since we didn't see [DONE]
+        assert drained is False
+        # Response should be closed
+        assert response.is_closed is True
+    finally:
+        # Restore original drain function
+        streaming_module.drain_async_iterator = original_drain
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("sync", [True, False], ids=["sync", "async"])
+async def test_drain_failure_after_done_preserves_result(sync: bool, client: OpenAI, async_client: AsyncOpenAI) -> None:
+    """Transport errors while draining after [DONE] must not fail an already-complete stream."""
+
+    def body() -> Iterator[bytes]:
+        yield b'data: {"foo":true}\n\n'
+        yield b"data: [DONE]\n\n"
+        raise httpx2.RemoteProtocolError("peer closed connection")
+
+    response = httpx2.Response(200, content=body() if sync else to_aiter(body()))
+
+    if sync:
+        stream: Stream[object] | AsyncStream[object] = Stream(cast_to=object, client=client, response=response)
+        chunks = list(stream)
+    else:
+        stream = AsyncStream(cast_to=object, client=async_client, response=response)
+        chunks = [chunk async for chunk in stream]
+
+    assert chunks == [{"foo": True}]
+    assert response.is_closed is True
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("sync", [True, False], ids=["sync", "async"])
+async def test_drain_decode_error_after_done_preserves_result(
+    sync: bool, client: OpenAI, async_client: AsyncOpenAI
+) -> None:
+    """Malformed trailing bytes after [DONE] must not fail an already-complete stream."""
+
+    def body() -> Iterator[bytes]:
+        yield b'data: {"foo":true}\n\n'
+        yield b"data: [DONE]\n\n"
+        # Truncated multi-byte UTF-8 sequence that the SSE decoder will reject.
+        yield b"data: \xff\n\n"
+
+    response = httpx2.Response(200, content=body() if sync else to_aiter(body()))
+
+    if sync:
+        stream: Stream[object] | AsyncStream[object] = Stream(cast_to=object, client=client, response=response)
+        chunks = list(stream)
+    else:
+        stream = AsyncStream(cast_to=object, client=async_client, response=response)
+        chunks = [chunk async for chunk in stream]
+
+    assert chunks == [{"foo": True}]
+    assert response.is_closed is True
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("sync", [True, False], ids=["sync", "async"])
+async def test_drain_is_bounded_and_doesnt_block_indefinitely(
+    sync: bool, client: OpenAI, async_client: AsyncOpenAI
+) -> None:
+    """Drain after [DONE] must timeout rather than wait for slow/infinite iterators."""
+
+    class SyncSlowIterator:
+        """Sync iterator that yields slowly after content to force timeout."""
+
+        def __init__(self, content: Iterator[bytes]) -> None:
+            self._content = content
+            self._content_exhausted = False
+
+        def __iter__(self) -> Iterator[bytes]:
+            return self
+
+        def __next__(self) -> bytes:
+            if not self._content_exhausted:
+                try:
+                    return next(self._content)
+                except StopIteration:
+                    self._content_exhausted = True
+            # After content is exhausted, yield slowly to force timeout during drain
+            import time
+
+            time.sleep(0.1)  # 100ms per item; 50ms drain timeout will hit after <1 item
+            return b": heartbeat\n\n"
+
+    class AsyncSlowIterator:
+        """Async iterator that yields slowly after content to force timeout."""
+
+        def __init__(self, content: AsyncIterator[bytes]) -> None:
+            self._content = content
+            self._content_exhausted = False
+
+        def __aiter__(self) -> AsyncIterator[bytes]:
+            return self
+
+        async def __anext__(self) -> bytes:
+            if not self._content_exhausted:
+                try:
+                    return await self._content.__anext__()
+                except StopAsyncIteration:
+                    self._content_exhausted = True
+            # After content is exhausted, yield slowly to force timeout during drain
+            import asyncio
+
+            await asyncio.sleep(0.1)  # 100ms per item; 50ms drain timeout will hit after <1 item
+            return b": heartbeat\n\n"
+
+    def body() -> Iterator[bytes]:
+        yield b'data: {"foo":true}\n\n'
+        yield b"data: [DONE]\n\n"
+
+    if sync:
+        slow_iter = SyncSlowIterator(body())
+        response = httpx2.Response(200, content=slow_iter)
+        stream: Stream[object] | AsyncStream[object] = Stream(cast_to=object, client=client, response=response)
+        chunks = list(stream)
+    else:
+        async_body = to_aiter(body())
+        slow_iter = AsyncSlowIterator(async_body)
+        response = httpx2.Response(200, content=slow_iter)
+        stream = AsyncStream(cast_to=object, client=async_client, response=response)
+        chunks = [chunk async for chunk in stream]
+
+    # Stream should complete quickly with just the first item, not wait for the slow iterator
+    assert chunks == [{"foo": True}]
+    assert response.is_closed is True
+
+
 async def to_aiter(iter: Iterator[bytes]) -> AsyncIterator[bytes]:
     for chunk in iter:
         yield chunk