From d4b47e7bb6a3611c5c344cda7582666276ec0a6c Mon Sep 17 00:00:00 2001 From: vrs-darkness Date: Thu, 6 Aug 2026 03:51:08 +0530 Subject: [PATCH 1/2] fix(streaming): ensure remaining body is consumed after [DONE] in Stream and AsyncStream - Added best-effort draining of remaining body bytes after receiving [DONE] to allow connection reuse. - Implemented error handling to prevent stream failures due to transport errors during draining. - Introduced tests to validate the behavior for both synchronous and asynchronous streams. --- src/openai/_streaming.py | 14 +++++++++- tests/test_streaming.py | 58 +++++++++++++++++++++++++++++++++++++--- 2 files changed, 68 insertions(+), 4 deletions(-) diff --git a/src/openai/_streaming.py b/src/openai/_streaming.py index 45c13cc11d..1e4ffbda4c 100644 --- a/src/openai/_streaming.py +++ b/src/openai/_streaming.py @@ -9,7 +9,7 @@ import httpx -from ._utils import is_mapping, extract_type_var_from_base +from ._utils import is_mapping, consume_sync_iterator, consume_async_iterator, extract_type_var_from_base from ._exceptions import APIError if TYPE_CHECKING: @@ -61,6 +61,12 @@ def __stream__(self) -> Iterator[_T]: try: for sse in iterator: if sse.data.startswith("[DONE]"): + # Best-effort drain so close() can return the connection to the pool. + # [DONE] is already terminal for callers; drain failures must not fail the stream. + try: + consume_sync_iterator(iterator) + except httpx.HTTPError: + pass break # we have to special case the Assistants `thread.` events since we won't have an "event" key in the data @@ -171,6 +177,12 @@ async def __stream__(self) -> AsyncIterator[_T]: try: async for sse in iterator: if sse.data.startswith("[DONE]"): + # Best-effort drain so aclose() can return the connection to the pool. + # [DONE] is already terminal for callers; drain failures must not fail the stream. + try: + await consume_async_iterator(iterator) + except httpx.HTTPError: + 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..a0961c3676 100644 --- a/tests/test_streaming.py +++ b/tests/test_streaming.py @@ -1,12 +1,12 @@ from __future__ import annotations -from typing import Iterator, AsyncIterator +from collections.abc import AsyncIterator, Iterator import httpx import pytest -from openai import OpenAI, AsyncOpenAI -from openai._streaming import Stream, AsyncStream, ServerSentEvent +from openai import AsyncOpenAI, OpenAI +from openai._streaming import AsyncStream, ServerSentEvent, Stream @pytest.mark.asyncio @@ -216,6 +216,58 @@ 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_body(sync: bool, client: OpenAI, async_client: AsyncOpenAI) -> None: + """After [DONE], remaining body bytes must be consumed so close() can reuse the connection.""" + exhausted = False + + def body() -> Iterator[bytes]: + nonlocal exhausted + yield b'data: {"foo":true}\n\n' + yield b"data: [DONE]\n\n" + yield b": trailing comment after done\n\n" + exhausted = True + + response = httpx.Response(200, content=body() if sync else to_aiter(body())) + + if sync: + stream: Stream[object] | AsyncStream[object] = Stream(cast_to=object, client=client, response=response) + chunks = list(stream) + else: + stream = AsyncStream(cast_to=object, client=async_client, response=response) + chunks = [chunk async for chunk in stream] + + assert chunks == [{"foo": True}] + assert exhausted is True + assert response.is_closed is True + + +@pytest.mark.asyncio +@pytest.mark.parametrize("sync", [True, False], ids=["sync", "async"]) +async def test_drain_failure_after_done_preserves_result( + sync: bool, client: OpenAI, async_client: AsyncOpenAI +) -> None: + """Transport errors while draining after [DONE] must not fail an already-complete stream.""" + + def body() -> Iterator[bytes]: + yield b'data: {"foo":true}\n\n' + yield b"data: [DONE]\n\n" + raise httpx.RemoteProtocolError("peer closed connection") + + response = httpx.Response(200, content=body() if sync else to_aiter(body())) + + if sync: + stream: Stream[object] | AsyncStream[object] = Stream(cast_to=object, client=client, response=response) + chunks = list(stream) + else: + stream = AsyncStream(cast_to=object, client=async_client, response=response) + chunks = [chunk async for chunk in stream] + + assert chunks == [{"foo": True}] + assert response.is_closed is True + + async def to_aiter(iter: Iterator[bytes]) -> AsyncIterator[bytes]: for chunk in iter: yield chunk From d59c4093884054397eed80e34023a03e4a4b62ba Mon Sep 17 00:00:00 2001 From: vrs-darkness Date: Thu, 6 Aug 2026 03:59:11 +0530 Subject: [PATCH 2/2] fix(streaming): handle UnicodeError during stream draining in Stream and AsyncStream - Updated error handling in both Stream and AsyncStream classes to catch UnicodeError in addition to HTTPError during the draining of the stream after receiving [DONE]. - Added a new test to ensure that malformed trailing bytes after [DONE] do not cause failures in already-complete streams for both synchronous and asynchronous scenarios. --- src/openai/_streaming.py | 4 ++-- tests/test_streaming.py | 26 ++++++++++++++++++++++++++ 2 files changed, 28 insertions(+), 2 deletions(-) diff --git a/src/openai/_streaming.py b/src/openai/_streaming.py index 1e4ffbda4c..51574e2c3e 100644 --- a/src/openai/_streaming.py +++ b/src/openai/_streaming.py @@ -65,7 +65,7 @@ def __stream__(self) -> Iterator[_T]: # [DONE] is already terminal for callers; drain failures must not fail the stream. try: consume_sync_iterator(iterator) - except httpx.HTTPError: + except (httpx.HTTPError, UnicodeError): pass break @@ -181,7 +181,7 @@ async def __stream__(self) -> AsyncIterator[_T]: # [DONE] is already terminal for callers; drain failures must not fail the stream. try: await consume_async_iterator(iterator) - except httpx.HTTPError: + except (httpx.HTTPError, UnicodeError): pass break diff --git a/tests/test_streaming.py b/tests/test_streaming.py index a0961c3676..dd4436b220 100644 --- a/tests/test_streaming.py +++ b/tests/test_streaming.py @@ -268,6 +268,32 @@ def body() -> Iterator[bytes]: assert response.is_closed is True +@pytest.mark.asyncio +@pytest.mark.parametrize("sync", [True, False], ids=["sync", "async"]) +async def test_drain_decode_error_after_done_preserves_result( + sync: bool, client: OpenAI, async_client: AsyncOpenAI +) -> None: + """Malformed trailing bytes after [DONE] must not fail an already-complete stream.""" + + def body() -> Iterator[bytes]: + yield b'data: {"foo":true}\n\n' + yield b"data: [DONE]\n\n" + # Truncated multi-byte UTF-8 sequence that the SSE decoder will reject. + yield b"data: \xff\n\n" + + response = httpx.Response(200, content=body() if sync else to_aiter(body())) + + if sync: + stream: Stream[object] | AsyncStream[object] = Stream(cast_to=object, client=client, response=response) + chunks = list(stream) + else: + stream = AsyncStream(cast_to=object, client=async_client, response=response) + chunks = [chunk async for chunk in stream] + + assert chunks == [{"foo": True}] + assert response.is_closed is True + + async def to_aiter(iter: Iterator[bytes]) -> AsyncIterator[bytes]: for chunk in iter: yield chunk