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