diff --git a/src/openai/_streaming.py b/src/openai/_streaming.py index 78e2d20aa7..8ee97e3458 100644 --- a/src/openai/_streaming.py +++ b/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 diff --git a/tests/test_streaming.py b/tests/test_streaming.py index ae6c0590f7..32b73ab12f 100644 --- a/tests/test_streaming.py +++ b/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