From 50b8df0fd28550a79183a195ff4f30570b0788bb Mon Sep 17 00:00:00 2001
From: Mahesh Sadupalli <mahesh.sadupalli@gmail.com>
Date: Thu, 23 Jul 2026 16:26:13 +0200
Subject: [PATCH] fix(streaming): drain remaining bytes after [DONE] before
 closing response

When Stream/__stream__ breaks on [DONE] it immediately runs the finally
block which calls response.close() / response.aclose(). If the HTTP/1.1
chunked terminator (0\r\n\r\n) has not yet been read by h11, their_state
is still SEND_RESPONSE at close() time, so httpcore takes the "destroy
connection" branch instead of returning it to the pool. This emits a
spurious TCP FIN and causes downstream proxies to log
downstream_remote_disconnect.

Regression introduced in 6132922c ("fix(client): close streams without
requiring full consumption") which removed the drain that 7e2b2544 had
originally added.

Fix: after detecting [DONE], exhaust the SSE iterator before breaking so
that h11 finishes parsing the chunked terminator. The drain only runs on
the clean [DONE] path; early consumer exit / error paths are unchanged.

Adds a regression test (sync + async) that verifies trailing bytes are
consumed before response.close() is called.

Fixes #3440
---
 src/openai/_streaming.py | 16 +++++++++++
 tests/test_streaming.py  | 61 ++++++++++++++++++++++++++++++++++++++++
 2 files changed, 77 insertions(+)

diff --git a/src/openai/_streaming.py b/src/openai/_streaming.py
index 45c13cc11d..39a132e840 100644
--- a/src/openai/_streaming.py
+++ b/src/openai/_streaming.py
@@ -61,6 +61,14 @@ def __stream__(self) -> Iterator[_T]:
         try:
             for sse in iterator:
                 if sse.data.startswith("[DONE]"):
+                    # Drain any remaining bytes so that h11 can parse the
+                    # HTTP/1.1 chunked terminator (`0\r\n\r\n`) before
+                    # response.close() is called.  Without this, h11's
+                    # their_state is still SEND_RESPONSE, so httpcore takes
+                    # the "destroy" branch instead of returning the connection
+                    # to the pool, causing a spurious TCP FIN.
+                    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 +179,14 @@ async def __stream__(self) -> AsyncIterator[_T]:
         try:
             async for sse in iterator:
                 if sse.data.startswith("[DONE]"):
+                    # Drain any remaining bytes so that h11 can parse the
+                    # HTTP/1.1 chunked terminator (`0\r\n\r\n`) before
+                    # response.aclose() is called.  Without this, h11's
+                    # their_state is still SEND_RESPONSE, so httpcore takes
+                    # the "destroy" branch instead of returning the connection
+                    # to the pool, causing a spurious TCP FIN.
+                    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
diff --git a/tests/test_streaming.py b/tests/test_streaming.py
index 04f8e51abd..895f67f065 100644
--- a/tests/test_streaming.py
+++ b/tests/test_streaming.py
@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 from typing import Iterator, AsyncIterator
+from unittest.mock import patch, MagicMock, AsyncMock
 
 import httpx
 import pytest
@@ -216,6 +217,66 @@ def body() -> Iterator[bytes]:
     assert sse.json() == {"content": "известни"}
 
 
+@pytest.mark.asyncio
+@pytest.mark.parametrize("sync", [True, False], ids=["sync", "async"])
+async def test_done_drains_remaining_bytes_before_close(
+    sync: bool, client: OpenAI, async_client: AsyncOpenAI
+) -> None:
+    """Regression test for #3440.
+
+    When Stream/__stream__ sees [DONE] it must drain the underlying byte
+    iterator before calling response.close() / response.aclose().  Without
+    the drain, h11's their_state is still SEND_RESPONSE at the time of
+    close(), so httpcore destroys the connection instead of returning it to
+    the pool, producing a spurious TCP FIN.
+    """
+    drain_calls: list[str] = []
+
+    def body() -> Iterator[bytes]:
+        yield b'data: {"id":"1"}\n\n'
+        yield b"data: [DONE]\n\n"
+        # bytes that arrive after [DONE] — the HTTP/1.1 chunked terminator
+        # would appear here in a real response
+        drain_calls.append("trailing-bytes-yielded")
+        yield b""
+
+    if sync:
+        response = httpx.Response(200, content=body())
+        close_calls: list[str] = []
+        original_close = response.close
+
+        def tracking_close() -> None:
+            close_calls.append("closed")
+            original_close()
+
+        response.close = tracking_close  # type: ignore[method-assign]
+        stream = Stream(cast_to=object, client=client, response=response)
+        list(stream.__stream__())
+
+        # The trailing bytes must have been consumed *before* close() was called.
+        assert drain_calls == ["trailing-bytes-yielded"], (
+            "trailing bytes were not drained before close()"
+        )
+    else:
+        async_body = to_aiter(body())
+        response = httpx.Response(200, content=async_body)
+        aclose_calls: list[str] = []
+        original_aclose = response.aclose
+
+        async def tracking_aclose() -> None:
+            aclose_calls.append("aclosed")
+            await original_aclose()
+
+        response.aclose = tracking_aclose  # type: ignore[method-assign]
+        stream = AsyncStream(cast_to=object, client=async_client, response=response)
+        async for _ in stream.__stream__():
+            pass
+
+        assert drain_calls == ["trailing-bytes-yielded"], (
+            "trailing bytes were not drained before aclose()"
+        )
+
+
 async def to_aiter(iter: Iterator[bytes]) -> AsyncIterator[bytes]:
     for chunk in iter:
         yield chunk