fix: drain remaining stream bytes after [DONE] event to prevent connection force-close

src/openai/_streaming.py

@@ -57,10 +57,12 @@ def __stream__(self) -> Iterator[_T]:
         response = self.response
         process_data = self._client._process_response_data
         iterator = self._iter_events()
+        terminated = False
 
         try:
             for sse in iterator:
                 if sse.data.startswith("[DONE]"):
+                    terminated = True
                     break
 
                 # we have to special case the Assistants `thread.` events since we won't have an "event" key in the data
@@ -106,8 +108,28 @@ def __stream__(self) -> Iterator[_T]:
                         response=response,
                     )
         finally:
-            # Ensure the response is closed even if the consumer doesn't read all data
-            response.close()
+            # Only drain when the stream terminated normally, i.e. we observed the
+            # `[DONE]` event and just need to consume the few remaining bytes (e.g. the
+            # HTTP/1.1 chunked terminator) so the connection can be returned to the pool.
+            #
+            # On premature termination -- the caller breaking out of iteration, an error,
+            # or cancellation -- draining would block until the server finishes sending,
+            # so close the response instead.
+            try:
+                if terminated:
+                    # Draining is best-effort cleanup for a stream that already completed.
+                    # If the connection drops before the trailing bytes arrive, the result
+                    # is still valid, so don't turn that into a failure for the caller.
+                    try:
+                        for _ in iterator:
+                            pass
+                    except httpx.HTTPError:
+                        pass
+            finally:
+                # The drain can still be interrupted by something that isn't an
+                # `httpx.HTTPError` -- e.g. the caller cancelling iteration while we
+                # wait on the trailing bytes. The connection must be released either way.
+                response.close()
 
     def __enter__(self) -> Self:
         return self
@@ -167,10 +189,12 @@ async def __stream__(self) -> AsyncIterator[_T]:
         response = self.response
         process_data = self._client._process_response_data
         iterator = self._iter_events()
+        terminated = False
 
         try:
             async for sse in iterator:
                 if sse.data.startswith("[DONE]"):
+                    terminated = True
                     break
 
                 # we have to special case the Assistants `thread.` events since we won't have an "event" key in the data
@@ -216,8 +240,30 @@ 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()
+            # Only drain when the stream terminated normally, i.e. we observed the
+            # `[DONE]` event and just need to consume the few remaining bytes (e.g. the
+            # HTTP/1.1 chunked terminator) so the connection can be returned to the pool.
+            #
+            # On premature termination -- the caller breaking out of iteration, an error,
+            # or cancellation -- draining would block until the server finishes sending,
+            # so close the response instead.
+            try:
+                if terminated:
+                    # Draining is best-effort cleanup for a stream that already completed.
+                    # If the connection drops before the trailing bytes arrive, the result
+                    # is still valid, so don't turn that into a failure for the caller.
+                    try:
+                        async for _ in iterator:
+                            pass
+                    except httpx.HTTPError:
+                        pass
+            finally:
+                # The drain can still be interrupted by something that isn't an
+                # `httpx.HTTPError` -- e.g. `asyncio.CancelledError` when the caller
+                # wraps consumption in a timeout shorter than the HTTPX read timeout.
+                # `CancelledError` is a `BaseException`, so it bypasses the handler
+                # above; the connection must be released either way.
+                await response.aclose()
 
     async def __aenter__(self) -> Self:
         return self

tests/test_streaming.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
-from typing import Iterator, AsyncIterator
+import asyncio
+from typing import Iterator, Generator, AsyncIterator, AsyncGenerator, cast
 
 import httpx
 import pytest
@@ -216,6 +217,153 @@ def body() -> Iterator[bytes]:
     assert sse.json() == {"content": "известни"}
 
 
+@pytest.mark.asyncio
+@pytest.mark.parametrize("sync", [True, False], ids=["sync", "async"])
+async def test_drains_remaining_bytes_after_done(
+    sync: bool,
+    client: OpenAI,
+    async_client: AsyncOpenAI,
+) -> None:
+    """After a normal `[DONE]` termination the trailing bytes are consumed so that
+    the underlying connection can be returned to the pool."""
+    consumed: list[bytes] = []
+
+    def body() -> Iterator[bytes]:
+        for chunk in [b'data: {"foo":true}\n\n', b"data: [DONE]\n\n", b"\n"]:
+            consumed.append(chunk)
+            yield chunk
+
+    stream = make_stream(content=body(), sync=sync, client=client, async_client=async_client)
+
+    items = await iter_all(stream)
+    assert len(items) == 1
+
+    # the whole body, including the bytes trailing `[DONE]`, was drained
+    assert consumed == [b'data: {"foo":true}\n\n', b"data: [DONE]\n\n", b"\n"]
+
+
+@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:
+    """A transport error while draining must not fail an already-completed stream.
+
+    The server sent a valid `[DONE]`, so the result is complete; a connection drop
+    while consuming the trailing bytes is cleanup noise and must still close the
+    response rather than propagating to the caller.
+    """
+
+    def body() -> Iterator[bytes]:
+        yield b'data: {"foo":true}\n\n'
+        yield b"data: [DONE]\n\n"
+        raise httpx.RemoteProtocolError("peer closed connection")
+
+    stream = make_stream(content=body(), sync=sync, client=client, async_client=async_client)
+
+    items = await iter_all(stream)
+    assert len(items) == 1
+
+    assert stream.response.is_closed
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("sync", [True, False], ids=["sync", "async"])
+async def test_drain_cancellation_after_done_still_closes_response(
+    sync: bool,
+    client: OpenAI,
+    async_client: AsyncOpenAI,
+) -> None:
+    """Cancellation while draining must still release the connection.
+
+    `[DONE]` has already been observed, so the drain is best-effort cleanup. If the
+    caller cancels while we wait on the trailing bytes -- e.g. an `asyncio` timeout
+    shorter than the HTTPX read timeout -- `CancelledError` is a `BaseException` and
+    so bypasses the `httpx.HTTPError` handler around the drain. The close has to sit
+    in a `finally` or the connection is leaked.
+    """
+
+    def body() -> Iterator[bytes]:
+        yield b'data: {"foo":true}\n\n'
+        yield b"data: [DONE]\n\n"
+        raise asyncio.CancelledError
+
+    stream = make_stream(content=body(), sync=sync, client=client, async_client=async_client)
+
+    with pytest.raises(asyncio.CancelledError):
+        await iter_all(stream)
+
+    assert stream.response.is_closed
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("sync", [True, False], ids=["sync", "async"])
+async def test_does_not_drain_on_premature_termination(
+    sync: bool,
+    client: OpenAI,
+    async_client: AsyncOpenAI,
+) -> None:
+    """A caller that stops before `[DONE]` must not drain the rest of the stream.
+
+    Draining here would block until the server finished sending, which for a
+    long-running or stalled completion defeats the point of breaking out early.
+    """
+    trailing = 100
+    consumed: list[bytes] = []
+
+    def body() -> Iterator[bytes]:
+        for chunk in [b'data: {"foo":true}\n\n', *([b'data: {"bar":true}\n\n'] * trailing), b"data: [DONE]\n\n"]:
+            consumed.append(chunk)
+            yield chunk
+
+    stream = make_stream(content=body(), sync=sync, client=client, async_client=async_client)
+
+    # consume a single event, then abandon iteration and let the stream be collected
+    await iter_next_item(stream)
+    await close_stream(stream)
+
+    # the remaining events were left on the wire rather than being drained
+    assert len(consumed) < trailing, f"stream was drained: consumed {len(consumed)} chunks"
+
+
+async def iter_all(stream: Stream[object] | AsyncStream[object]) -> list[object]:
+    if isinstance(stream, AsyncStream):
+        return [item async for item in stream]
+
+    return list(stream)
+
+
+async def iter_next_item(stream: Stream[object] | AsyncStream[object]) -> object:
+    if isinstance(stream, AsyncStream):
+        return await stream.__anext__()
+
+    return next(iter(stream))
+
+
+async def close_stream(stream: Stream[object] | AsyncStream[object]) -> None:
+    """Close the underlying generator, mirroring what happens when an abandoned
+    stream object is garbage collected."""
+    if isinstance(stream, AsyncStream):
+        await cast("AsyncGenerator[object, None]", stream._iterator).aclose()
+    else:
+        cast("Generator[object, None, None]", stream._iterator).close()
+
+
+def make_stream(
+    content: Iterator[bytes],
+    *,
+    sync: bool,
+    client: OpenAI,
+    async_client: AsyncOpenAI,
+) -> Stream[object] | AsyncStream[object]:
+    if sync:
+        return Stream(cast_to=object, client=client, response=httpx.Response(200, content=content))
+
+    return AsyncStream(cast_to=object, client=async_client, response=httpx.Response(200, content=to_aiter(content)))
+
+
 async def to_aiter(iter: Iterator[bytes]) -> AsyncIterator[bytes]:
     for chunk in iter:
         yield chunk