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

src/openai/_streaming.py

@@ -62,6 +62,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
@@ -172,6 +185,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

@@ -246,3 +246,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