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

src/openai/_streaming.py

@@ -19,13 +19,22 @@
 
 _T = TypeVar("_T")
 
+# After a normal `[DONE]` termination, only this many bytes of the response are
+# consumed while draining. In the healthy case the trailing bytes are just
+# transport framing (e.g. the HTTP/1.1 chunked terminator) -- a handful of bytes
+# at most -- so this is generous headroom. It exists so that a server or proxy
+# that leaves the connection open (or keeps sending data) after `[DONE]` can't
+# turn "drain the last few bytes" into an unbounded read.
+_MAX_POST_DONE_DRAIN_BYTES = 64 * 1024
+
 
 class Stream(Generic[_T]):
     """Provides the core interface to iterate over a synchronous stream response."""
 
     response: httpx.Response
     _options: Optional[FinalRequestOptions] = None
     _decoder: SSEBytesDecoder
+    _terminated: bool = False
 
     def __init__(
         self,
@@ -50,17 +59,35 @@ def __iter__(self) -> Iterator[_T]:
             yield item
 
     def _iter_events(self) -> Iterator[ServerSentEvent]:
-        yield from self._decoder.iter_bytes(self.response.iter_bytes())
+        yield from self._decoder.iter_bytes(self._iter_raw_bytes())
+
+    def _iter_raw_bytes(self) -> Iterator[bytes]:
+        """Wraps `response.iter_bytes()` so the post-`[DONE]` drain is bounded.
+
+        Once `__stream__` sets `self._terminated` (having observed `[DONE]`), this
+        stops yielding new bytes after `_MAX_POST_DONE_DRAIN_BYTES` have been
+        consumed, so a connection that stays open -- or keeps sending data -- after
+        the stream has logically completed can't stall the drain indefinitely.
+        """
+        drained = 0
+        for chunk in self.response.iter_bytes():
+            yield chunk
+            if self._terminated:
+                drained += len(chunk)
+                if drained >= _MAX_POST_DONE_DRAIN_BYTES:
+                    return
 
     def __stream__(self) -> Iterator[_T]:
         cast_to = cast(Any, self._cast_to)
         response = self.response
         process_data = self._client._process_response_data
+        self._terminated = False
         iterator = self._iter_events()
 
         try:
             for sse in iterator:
                 if sse.data.startswith("[DONE]"):
+                    self._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 +133,31 @@ 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.
+            # `_iter_raw_bytes()` bounds how much of that drain we're willing to do, so a
+            # connection that stays open (or keeps sending data) after `[DONE]` can't turn
+            # this into an unbounded read.
+            #
+            # 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 self._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
@@ -135,6 +185,7 @@ class AsyncStream(Generic[_T]):
     response: httpx.Response
     _options: Optional[FinalRequestOptions] = None
     _decoder: SSEDecoder | SSEBytesDecoder
+    _terminated: bool = False
 
     def __init__(
         self,
@@ -159,18 +210,36 @@ 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()):
+        async for sse in self._decoder.aiter_bytes(self._aiter_raw_bytes()):
             yield sse
 
+    async def _aiter_raw_bytes(self) -> AsyncIterator[bytes]:
+        """Wraps `response.aiter_bytes()` so the post-`[DONE]` drain is bounded.
+
+        Once `__stream__` sets `self._terminated` (having observed `[DONE]`), this
+        stops yielding new bytes after `_MAX_POST_DONE_DRAIN_BYTES` have been
+        consumed, so a connection that stays open -- or keeps sending data -- after
+        the stream has logically completed can't stall the drain indefinitely.
+        """
+        drained = 0
+        async for chunk in self.response.aiter_bytes():
+            yield chunk
+            if self._terminated:
+                drained += len(chunk)
+                if drained >= _MAX_POST_DONE_DRAIN_BYTES:
+                    return
+
     async def __stream__(self) -> AsyncIterator[_T]:
         cast_to = cast(Any, self._cast_to)
         response = self.response
         process_data = self._client._process_response_data
+        self._terminated = False
         iterator = self._iter_events()
 
         try:
             async for sse in iterator:
                 if sse.data.startswith("[DONE]"):
+                    self._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 +285,33 @@ 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.
+            # `_aiter_raw_bytes()` bounds how much of that drain we're willing to do, so
+            # a connection that stays open (or keeps sending data) after `[DONE]` can't
+            # turn this into an unbounded read.
+            #
+            # 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 self._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,12 +1,13 @@
 from __future__ import annotations
 
-from typing import Iterator, AsyncIterator
+import asyncio
+from typing import Iterator, Generator, AsyncIterator, AsyncGenerator, cast
 
 import httpx
 import pytest
 
 from openai import OpenAI, AsyncOpenAI
-from openai._streaming import Stream, AsyncStream, ServerSentEvent
+from openai._streaming import _MAX_POST_DONE_DRAIN_BYTES, Stream, AsyncStream, ServerSentEvent
 
 
 @pytest.mark.asyncio
@@ -216,6 +217,196 @@ 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_post_done_drain_is_bounded(
+    sync: bool,
+    client: OpenAI,
+    async_client: AsyncOpenAI,
+) -> None:
+    """The post-`[DONE]` drain must not turn into an unbounded read.
+
+    A server or a proxy in between can send `[DONE]` but leave the HTTP response
+    open (or keep trickling bytes on it) instead of closing the connection. The
+    stream has already logically completed at that point, so draining should stop
+    after a bounded amount of data rather than consuming everything the other end
+    is willing to send.
+    """
+    # Comment lines never decode into an `SSE` event, so none of these chunks make
+    # it back to the `for _ in iterator: pass` drain loop as a yielded event --
+    # the only thing that can bound consumption here is the underlying byte read
+    # itself, not anything at the parsed-event level.
+    trailing_chunk = b": padding to simulate a proxy holding the connection open\n\n"
+    # comfortably more than the drain is allowed to consume
+    num_trailing_chunks = (_MAX_POST_DONE_DRAIN_BYTES // len(trailing_chunk)) * 4
+
+    consumed = 0
+
+    def body() -> Iterator[bytes]:
+        nonlocal consumed
+        yield b'data: {"foo":true}\n\n'
+        yield b"data: [DONE]\n\n"
+        for _ in range(num_trailing_chunks):
+            consumed += len(trailing_chunk)
+            yield trailing_chunk
+
+    stream = make_stream(content=body(), sync=sync, client=client, async_client=async_client)
+
+    items = await iter_all(stream)
+    assert len(items) == 1
+
+    # the drain stopped well short of consuming every trailing chunk the "server" sent
+    assert consumed < num_trailing_chunks * len(trailing_chunk)
+    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