fix: drain remaining bytes after [DONE] before closing response (#3440)

When the SSE decoder encounters [DONE], it breaks out of the event
loop immediately. However, the underlying httpx response iterator may
not have reached EOF yet — the chunked terminator (0\r\n\r\n) can still
be in flight. Calling response.close() at this point sends a TCP FIN
while the server is still transmitting, which prevents h11 from
advancing to the DONE state and causes connection pool degradation and
proxy errors.

The fix drains remaining events from the existing iterator (not a new
response.iter_bytes() call, which would raise StreamConsumed) after
[DONE] is encountered, in both sync and async paths. This allows the
response to reach EOF naturally before close.

Regression tests verify that trailing events after [DONE] are consumed
and the response is fully closed in both sync and async modes.

Closes #3440

Author: rkfshakti

src/openai/_streaming.py

@@ -68,6 +68,19 @@ def __stream__(self) -> Iterator[_T]:
         try:
             for sse in iterator:
                 if sse.data.startswith("[DONE]"):
+                    # Drain remaining events from the existing iterator so the
+                    # underlying response.iter_bytes() reaches EOF, allowing
+                    # h11 to advance to DONE state before close. Without this,
+                    # response.close() sends TCP FIN while the chunked terminator
+                    # (0\r\n\r\n) is still in flight, causing connection pool
+                    # degradation and proxy errors. (#3440)
+                    #
+                    # We must drain through `iterator` (not start a new
+                    # `self.response.iter_bytes()`) because httpx only allows
+                    # one active iterator at a time — a second call raises
+                    # `httpx.StreamConsumed`.
+                    for _ in iterator:
+                        pass
                     break
 
                 # we have to special case the Assistants `thread.` events since we won't have an "event" key in the data
@@ -183,6 +196,19 @@ async def __stream__(self) -> AsyncIterator[_T]:
         try:
             async for sse in iterator:
                 if sse.data.startswith("[DONE]"):
+                    # Drain remaining events from the existing iterator so the
+                    # underlying response.aiter_bytes() reaches EOF, allowing
+                    # h11 to advance to DONE state before close. Without this,
+                    # response.aclose() sends TCP FIN while the chunked terminator
+                    # (0\r\n\r\n) is still in flight, causing connection pool
+                    # degradation and proxy errors. (#3440)
+                    #
+                    # We must drain through `iterator` (not start a new
+                    # `self.response.aiter_bytes()`) because httpx only allows
+                    # one active iterator at a time — a second call raises
+                    # `httpx.StreamConsumed`.
+                    async for _ in iterator:
+                        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

@@ -438,3 +438,37 @@ def make_event_iterator(
     return AsyncStream(
         cast_to=object, client=async_client, response=httpx2.Response(200, content=to_aiter(content))
     )._iter_events()
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("sync", [True, False], ids=["sync", "async"])
+async def test_drain_after_done_consumes_trailing_events(sync: bool, client: OpenAI, async_client: AsyncOpenAI) -> None:
+    """After [DONE], the stream should drain remaining events from the iterator
+    so the underlying response reaches EOF. Regression test for #3440."""
+
+    def body() -> Iterator[bytes]:
+        yield b"event: completion\n"
+        yield b'data: {"foo":true}\n'
+        yield b"\n"
+        yield b"data: [DONE]\n"
+        yield b"\n"
+        # Trailing event after [DONE] — should be consumed by the drain.
+        yield b"event: trailing\n"
+        yield b'data: {"bar":false}\n'
+        yield b"\n"
+
+    if sync:
+        response = httpx2.Response(200, content=body())
+        stream = Stream(cast_to=object, client=client, response=response)
+        results: list[object] = list(stream)
+        assert len(results) == 1
+        assert results[0] == {"foo": True}
+        # The response should be fully consumed (not just half-read).
+        assert response.is_closed
+    else:
+        response = httpx2.Response(200, content=to_aiter(body()))
+        stream = AsyncStream(cast_to=object, client=async_client, response=response)
+        results = [item async for item in stream]  # type: ignore[reportUnknownVariableType]
+        assert len(results) == 1
+        assert results[0] == {"foo": True}
+        assert response.is_closed