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

src/openai/_streaming.py

@@ -61,6 +61,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
@@ -171,6 +184,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

@@ -216,6 +216,57 @@ def body() -> Iterator[bytes]:
     assert sse.json() == {"content": "известни"}
 
 
+@pytest.mark.asyncio
+@pytest.mark.parametrize("sync", [True, False], ids=["sync", "async"])
+async def test_stream_drains_after_done(
+    sync: bool,
+    client: OpenAI,
+    async_client: AsyncOpenAI,
+) -> None:
+    """Regression test for #3440: after [DONE], remaining events should be drained.
+
+    The fix drains the existing iterator (not a new response.iter_bytes() call)
+    so that httpx doesn't raise StreamConsumed and the response reaches EOF
+    before close, preventing premature TCP FIN.
+
+    We track whether the body generator reached its final sentinel to verify
+    that the drain actually consumed the trailing event — not just that the
+    response was closed (which happens in the finally block regardless).
+    """
+
+    consumed_after_done: list[bool] = []
+
+    def body() -> Iterator[bytes]:
+        yield b'data: {"foo":true}\n'
+        yield b"\n"
+        yield b"data: [DONE]\n"
+        yield b"\n"
+        # Extra data after [DONE] that should be consumed by the drain
+        yield b'data: {"trailing":true}\n'
+        yield b"\n"
+        # Sentinel: only reached if the drain consumed all trailing events
+        consumed_after_done.append(True)
+
+    if sync:
+        response = httpx.Response(200, content=body())
+        stream = Stream(cast_to=object, client=client, response=response)
+        # Consume the full stream — the drain should consume the trailing event
+        for _ in stream:
+            pass
+        assert response.is_closed
+    else:
+        response = httpx.Response(200, content=to_aiter(body()))
+        stream = AsyncStream(cast_to=object, client=async_client, response=response)
+        async for _ in stream:
+            pass
+        assert response.is_closed
+
+    # The sentinel is only appended if the body generator was fully consumed.
+    # Without the drain loop, the generator is garbage-collected when the
+    # response closes, so this assertion fails — proving the drain works.
+    assert consumed_after_done == [True]
+
+
 async def to_aiter(iter: Iterator[bytes]) -> AsyncIterator[bytes]:
     for chunk in iter:
         yield chunk