diff --git a/src/openai/_streaming.py b/src/openai/_streaming.py index 45c13cc11d..51574e2c3e 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, UnicodeError): + 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, UnicodeError): + 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..dd4436b220 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,84 @@ 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 + + +@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