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

src/openai/_streaming.py

@@ -9,7 +9,7 @@
 
 import httpx
 
-from ._utils import is_mapping, extract_type_var_from_base
+from ._utils import is_mapping, consume_sync_iterator, consume_async_iterator, extract_type_var_from_base
 from ._exceptions import APIError
 
 if TYPE_CHECKING:
@@ -61,6 +61,12 @@ def __stream__(self) -> Iterator[_T]:
         try:
             for sse in iterator:
                 if sse.data.startswith("[DONE]"):
+                    # Best-effort drain so close() can return the connection to the pool.
+                    # [DONE] is already terminal for callers; drain failures must not fail the stream.
+                    try:
+                        consume_sync_iterator(iterator)
+                    except (httpx.HTTPError, UnicodeError):
+                        pass
                     break
 
                 # we have to special case the Assistants `thread.` events since we won't have an "event" key in the data
@@ -171,6 +177,12 @@ async def __stream__(self) -> AsyncIterator[_T]:
         try:
             async for sse in iterator:
                 if sse.data.startswith("[DONE]"):
+                    # Best-effort drain so aclose() can return the connection to the pool.
+                    # [DONE] is already terminal for callers; drain failures must not fail the stream.
+                    try:
+                        await consume_async_iterator(iterator)
+                    except (httpx.HTTPError, UnicodeError):
+                        pass
                     break
 
                 # we have to special case the Assistants `thread.` events since we won't have an "event" key in the data

tests/test_streaming.py

@@ -1,12 +1,12 @@
 from __future__ import annotations
 
-from typing import Iterator, AsyncIterator
+from collections.abc import AsyncIterator, Iterator
 
 import httpx
 import pytest
 
-from openai import OpenAI, AsyncOpenAI
-from openai._streaming import Stream, AsyncStream, ServerSentEvent
+from openai import AsyncOpenAI, OpenAI
+from openai._streaming import AsyncStream, ServerSentEvent, Stream
 
 
 @pytest.mark.asyncio
@@ -216,6 +216,84 @@ def body() -> Iterator[bytes]:
     assert sse.json() == {"content": "известни"}
 
 
+@pytest.mark.asyncio
+@pytest.mark.parametrize("sync", [True, False], ids=["sync", "async"])
+async def test_done_drains_remaining_body(sync: bool, client: OpenAI, async_client: AsyncOpenAI) -> None:
+    """After [DONE], remaining body bytes must be consumed so close() can reuse the connection."""
+    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
+
+    response = httpx.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 exhausted is True
+    assert response.is_closed is True
+
+
+@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 httpx.RemoteProtocolError("peer closed connection")
+
+    response = httpx.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 = httpx.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
+
+
 async def to_aiter(iter: Iterator[bytes]) -> AsyncIterator[bytes]:
     for chunk in iter:
         yield chunk