diff --git a/src/openai/_streaming.py b/src/openai/_streaming.py index 45c13cc11d..0a7a38af63 100644 --- a/src/openai/_streaming.py +++ b/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 diff --git a/tests/test_streaming.py b/tests/test_streaming.py index 04f8e51abd..de7be5ec14 100644 --- a/tests/test_streaming.py +++ b/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