From d4b47e7bb6a3611c5c344cda7582666276ec0a6c Mon Sep 17 00:00:00 2001 From: vrs-darkness Date: Thu, 6 Aug 2026 03:51:08 +0530 Subject: [PATCH 01/21] 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 02/21] 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 From 114ec5741f7457e8bb9725058e5fe86e3965b382 Mon Sep 17 00:00:00 2001 From: vrs-darkness Date: Fri, 14 Aug 2026 18:47:47 +0530 Subject: [PATCH 03/21] refactor(streaming): replace consume functions with drain functions for bounded iterator handling - Updated the streaming logic to use `drain_sync_iterator` and `drain_async_iterator` instead of `consume_sync_iterator` and `consume_async_iterator`. - Enhanced comments to clarify the purpose of bounded draining after stream termination signals. - Added new drain functions to handle a limited number of items from iterators, preventing indefinite blocking. This change improves resource management and connection reuse after stream completion. --- src/openai/_streaming.py | 14 +++--- src/openai/_utils/__init__.py | 7 ++- src/openai/_utils/_streams.py | 26 ++++++++++ tests/test_streaming.py | 93 +++++++++++++++++++++++++++++++---- 4 files changed, 122 insertions(+), 18 deletions(-) diff --git a/src/openai/_streaming.py b/src/openai/_streaming.py index 5dec0a9279..f4ebe00a36 100644 --- a/src/openai/_streaming.py +++ b/src/openai/_streaming.py @@ -9,7 +9,7 @@ import httpx2 -from ._utils import is_mapping, consume_sync_iterator, consume_async_iterator, extract_type_var_from_base +from ._utils import is_mapping, drain_sync_iterator, drain_async_iterator, extract_type_var_from_base from ._exceptions import APIError if TYPE_CHECKING: @@ -61,11 +61,11 @@ 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. + # Best-effort bounded 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): + drain_sync_iterator(iterator) + except (httpx2.HTTPError, UnicodeError): pass break @@ -177,11 +177,11 @@ 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. + # Best-effort bounded 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): + await drain_async_iterator(iterator) + except (httpx2.HTTPError, UnicodeError): pass break diff --git a/src/openai/_utils/__init__.py b/src/openai/_utils/__init__.py index bbd79691fa..a91fc2f6c1 100644 --- a/src/openai/_utils/__init__.py +++ b/src/openai/_utils/__init__.py @@ -52,7 +52,12 @@ strip_annotated_type as strip_annotated_type, extract_type_var_from_base as extract_type_var_from_base, ) -from ._streams import consume_sync_iterator as consume_sync_iterator, consume_async_iterator as consume_async_iterator +from ._streams import ( + drain_sync_iterator as drain_sync_iterator, + drain_async_iterator as drain_async_iterator, + consume_sync_iterator as consume_sync_iterator, + consume_async_iterator as consume_async_iterator, +) from ._transform import ( PropertyInfo as PropertyInfo, transform as transform, diff --git a/src/openai/_utils/_streams.py b/src/openai/_utils/_streams.py index f4a0208f01..08617c8723 100644 --- a/src/openai/_utils/_streams.py +++ b/src/openai/_utils/_streams.py @@ -10,3 +10,29 @@ def consume_sync_iterator(iterator: Iterator[Any]) -> None: async def consume_async_iterator(iterator: AsyncIterator[Any]) -> None: async for _ in iterator: ... + + +def drain_sync_iterator(iterator: Iterator[Any], max_items: int = 16) -> None: + """Drain a bounded number of items from an iterator without blocking indefinitely. + + Used after stream termination signals like [DONE] to attempt connection reuse + without waiting for the entire response body. + """ + for _ in range(max_items): + try: + next(iterator) + except StopIteration: + break + + +async def drain_async_iterator(iterator: AsyncIterator[Any], max_items: int = 16) -> None: + """Drain a bounded number of items from an async iterator without blocking indefinitely. + + Used after stream termination signals like [DONE] to attempt connection reuse + without waiting for the entire response body. + """ + for _ in range(max_items): + try: + await iterator.__anext__() + except StopAsyncIteration: + break diff --git a/tests/test_streaming.py b/tests/test_streaming.py index 9a01b0a46b..30bf4ad4d0 100644 --- a/tests/test_streaming.py +++ b/tests/test_streaming.py @@ -1,12 +1,12 @@ from __future__ import annotations -from collections.abc import AsyncIterator, Iterator +from collections.abc import Iterator, AsyncIterator import httpx2 import pytest -from openai import AsyncOpenAI, OpenAI -from openai._streaming import AsyncStream, ServerSentEvent, Stream +from openai import OpenAI, AsyncOpenAI +from openai._streaming import Stream, AsyncStream, ServerSentEvent @pytest.mark.asyncio @@ -229,7 +229,7 @@ def body() -> Iterator[bytes]: yield b": trailing comment after done\n\n" exhausted = True - response = httpx.Response(200, content=body() if sync else to_aiter(body())) + response = httpx2.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) @@ -245,17 +245,15 @@ def body() -> Iterator[bytes]: @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: +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") + raise httpx2.RemoteProtocolError("peer closed connection") - response = httpx.Response(200, content=body() if sync else to_aiter(body())) + response = httpx2.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) @@ -281,15 +279,90 @@ def body() -> Iterator[bytes]: # 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())) + response = httpx2.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_is_bounded_and_doesnt_block_indefinitely( + sync: bool, client: OpenAI, async_client: AsyncOpenAI +) -> None: + """Drain after [DONE] must be bounded; trailing iterator that never closes must not block completion.""" + + class SyncInfiniteIterator: + """Sync iterator that yields indefinitely after content; simulates unbounded heartbeats.""" + + def __init__(self, content: Iterator[bytes]) -> None: + self._content = content + self._content_exhausted = False + self._yield_count = 0 + + def __iter__(self) -> Iterator[bytes]: + return self + + def __next__(self) -> bytes: + if not self._content_exhausted: + try: + return next(self._content) + except StopIteration: + self._content_exhausted = True + # After content is exhausted, yield indefinitely (simulating held-open connection) + if self._yield_count > 1000: # Safety limit to prevent infinite test loops + raise StopIteration + self._yield_count += 1 + return b": heartbeat\n\n" + + class AsyncInfiniteIterator: + """Async iterator that yields indefinitely after content; simulates unbounded heartbeats.""" + + def __init__(self, content: AsyncIterator[bytes]) -> None: + self._content = content + self._content_exhausted = False + self._yield_count = 0 + + def __aiter__(self) -> AsyncIterator[bytes]: + return self + + async def __anext__(self) -> bytes: + if not self._content_exhausted: + try: + return await self._content.__anext__() + except StopAsyncIteration: + self._content_exhausted = True + # After content is exhausted, yield indefinitely (simulating held-open connection) + if self._yield_count > 1000: # Safety limit to prevent infinite test loops + raise StopAsyncIteration + self._yield_count += 1 + return b": heartbeat\n\n" + + def body() -> Iterator[bytes]: + yield b'data: {"foo":true}\n\n' + yield b"data: [DONE]\n\n" + # More events after [DONE] that should not block completion if sync: + infinite_iter = SyncInfiniteIterator(body()) + response = httpx2.Response(200, content=infinite_iter) stream: Stream[object] | AsyncStream[object] = Stream(cast_to=object, client=client, response=response) chunks = list(stream) else: + async_body = to_aiter(body()) + infinite_iter = AsyncInfiniteIterator(async_body) + response = httpx2.Response(200, content=infinite_iter) stream = AsyncStream(cast_to=object, client=async_client, response=response) chunks = [chunk async for chunk in stream] + # Stream should complete with just the first item, not block waiting for the infinite iterator assert chunks == [{"foo": True}] assert response.is_closed is True From 8a8d02b20e670d59a84f66332657434a4fa89621 Mon Sep 17 00:00:00 2001 From: vrs-darkness Date: Fri, 14 Aug 2026 18:56:22 +0530 Subject: [PATCH 04/21] refactor(streaming): update iterator handling for byte streams - Modified the streaming logic to utilize `iter_bytes` and `aiter_bytes` for both synchronous and asynchronous streams. - Enhanced comments to clarify the purpose of draining raw bytes after stream termination signals, ensuring proper resource management and preventing unbounded waits. This change improves the handling of byte streams in the streaming classes. --- src/openai/_streaming.py | 16 ++++++++++------ 1 file changed, 10 insertions(+), 6 deletions(-) diff --git a/src/openai/_streaming.py b/src/openai/_streaming.py index f4ebe00a36..dc23cb61fc 100644 --- a/src/openai/_streaming.py +++ b/src/openai/_streaming.py @@ -56,15 +56,17 @@ def __stream__(self) -> Iterator[_T]: cast_to = cast(Any, self._cast_to) response = self.response process_data = self._client._process_response_data - iterator = self._iter_events() + byte_iterator = response.iter_bytes() + iterator = self._decoder.iter_bytes(byte_iterator) try: for sse in iterator: if sse.data.startswith("[DONE]"): - # Best-effort bounded drain so close() can return the connection to the pool. + # Best-effort bounded drain of raw bytes so close() can return the connection to the pool. # [DONE] is already terminal for callers; drain failures must not fail the stream. + # Drain raw bytes (not decoded events) so heartbeat comments don't cause unbounded waits. try: - drain_sync_iterator(iterator) + drain_sync_iterator(byte_iterator) except (httpx2.HTTPError, UnicodeError): pass break @@ -172,15 +174,17 @@ async def __stream__(self) -> AsyncIterator[_T]: cast_to = cast(Any, self._cast_to) response = self.response process_data = self._client._process_response_data - iterator = self._iter_events() + byte_iterator = response.aiter_bytes() + iterator = self._decoder.aiter_bytes(byte_iterator) try: async for sse in iterator: if sse.data.startswith("[DONE]"): - # Best-effort bounded drain so aclose() can return the connection to the pool. + # Best-effort bounded drain of raw bytes so aclose() can return the connection to the pool. # [DONE] is already terminal for callers; drain failures must not fail the stream. + # Drain raw bytes (not decoded events) so heartbeat comments don't cause unbounded waits. try: - await drain_async_iterator(iterator) + await drain_async_iterator(byte_iterator) except (httpx2.HTTPError, UnicodeError): pass break From 7e528cb277eb144754ce52c13e4680ad1f803f26 Mon Sep 17 00:00:00 2001 From: vrs-darkness Date: Fri, 14 Aug 2026 19:27:08 +0530 Subject: [PATCH 05/21] refactor(streaming): enhance iterator draining with timeout handling - Updated `drain_sync_iterator` and `drain_async_iterator` to implement timeout-based draining of remaining items from iterators, improving connection reuse after stream termination. - Enhanced comments to clarify the purpose of bounded draining and ensure proper resource management without indefinite blocking. This change optimizes the handling of byte streams in both synchronous and asynchronous contexts. --- src/openai/_streaming.py | 32 ++++++++++---------- src/openai/_utils/_streams.py | 55 ++++++++++++++++++++++------------- 2 files changed, 50 insertions(+), 37 deletions(-) diff --git a/src/openai/_streaming.py b/src/openai/_streaming.py index dc23cb61fc..67e439333f 100644 --- a/src/openai/_streaming.py +++ b/src/openai/_streaming.py @@ -62,13 +62,6 @@ def __stream__(self) -> Iterator[_T]: try: for sse in iterator: if sse.data.startswith("[DONE]"): - # Best-effort bounded drain of raw bytes so close() can return the connection to the pool. - # [DONE] is already terminal for callers; drain failures must not fail the stream. - # Drain raw bytes (not decoded events) so heartbeat comments don't cause unbounded waits. - try: - drain_sync_iterator(byte_iterator) - except (httpx2.HTTPError, UnicodeError): - pass break # we have to special case the Assistants `thread.` events since we won't have an "event" key in the data @@ -114,7 +107,14 @@ def __stream__(self) -> Iterator[_T]: response=response, ) finally: - # Ensure the response is closed even if the consumer doesn't read all data + # Best-effort timeout-bounded drain of raw bytes to enable connection reuse. + # [DONE] is already terminal for callers; drain failures must not fail the stream. + # Drain raw bytes (not decoded events) so heartbeat comments don't cause unbounded waits. + # Drain in finally so [DONE] doesn't block the caller. + try: + drain_sync_iterator(byte_iterator, timeout_ms=50) + except (httpx2.HTTPError, UnicodeError): + pass response.close() def __enter__(self) -> Self: @@ -180,13 +180,6 @@ async def __stream__(self) -> AsyncIterator[_T]: try: async for sse in iterator: if sse.data.startswith("[DONE]"): - # Best-effort bounded drain of raw bytes so aclose() can return the connection to the pool. - # [DONE] is already terminal for callers; drain failures must not fail the stream. - # Drain raw bytes (not decoded events) so heartbeat comments don't cause unbounded waits. - try: - await drain_async_iterator(byte_iterator) - except (httpx2.HTTPError, UnicodeError): - pass break # we have to special case the Assistants `thread.` events since we won't have an "event" key in the data @@ -232,7 +225,14 @@ async def __stream__(self) -> AsyncIterator[_T]: response=response, ) finally: - # Ensure the response is closed even if the consumer doesn't read all data + # Best-effort timeout-bounded drain of raw bytes to enable connection reuse. + # [DONE] is already terminal for callers; drain failures must not fail the stream. + # Drain raw bytes (not decoded events) so heartbeat comments don't cause unbounded waits. + # Drain in finally so [DONE] doesn't block the caller. + try: + await drain_async_iterator(byte_iterator, timeout_ms=50) + except (httpx2.HTTPError, UnicodeError): + pass await response.aclose() async def __aenter__(self) -> Self: diff --git a/src/openai/_utils/_streams.py b/src/openai/_utils/_streams.py index 08617c8723..65408c3c9d 100644 --- a/src/openai/_utils/_streams.py +++ b/src/openai/_utils/_streams.py @@ -1,3 +1,4 @@ +import asyncio from typing import Any from typing_extensions import Iterator, AsyncIterator @@ -12,27 +13,39 @@ async def consume_async_iterator(iterator: AsyncIterator[Any]) -> None: ... -def drain_sync_iterator(iterator: Iterator[Any], max_items: int = 16) -> None: - """Drain a bounded number of items from an iterator without blocking indefinitely. +def drain_sync_iterator(iterator: Iterator[Any], timeout_ms: int = 50) -> None: + """Drain trailing bytes from iterator with bounded timeout. - Used after stream termination signals like [DONE] to attempt connection reuse - without waiting for the entire response body. + Attempts to drain all remaining items from iterator to enable connection + reuse, but gives up after timeout_ms to avoid indefinite blocking when + server holds connection open after [DONE]. """ - for _ in range(max_items): - try: - next(iterator) - except StopIteration: - break - - -async def drain_async_iterator(iterator: AsyncIterator[Any], max_items: int = 16) -> None: - """Drain a bounded number of items from an async iterator without blocking indefinitely. - - Used after stream termination signals like [DONE] to attempt connection reuse - without waiting for the entire response body. + import time + deadline = time.monotonic() + (timeout_ms / 1000.0) + try: + while time.monotonic() < deadline: + try: + next(iterator) + except StopIteration: + return + except Exception: + pass + + +async def drain_async_iterator(iterator: AsyncIterator[Any], timeout_ms: int = 50) -> None: + """Drain trailing bytes from async iterator with bounded timeout. + + Attempts to drain all remaining items from iterator to enable connection + reuse, but gives up after timeout_ms to avoid indefinite blocking when + server holds connection open after [DONE]. """ - for _ in range(max_items): - try: - await iterator.__anext__() - except StopAsyncIteration: - break + try: + while True: + try: + await asyncio.wait_for(iterator.__anext__(), timeout=timeout_ms / 1000.0) + except StopAsyncIteration: + return + except asyncio.TimeoutError: + pass + except Exception: + pass From 0316faab49d3b2657df0f6a6149e8fc175b33bc8 Mon Sep 17 00:00:00 2001 From: vrs-darkness Date: Fri, 14 Aug 2026 19:27:47 +0530 Subject: [PATCH 06/21] fix(streaming): add missing newline for code clarity - Added a newline in the `drain_sync_iterator` function to improve code readability and maintain consistency in formatting. This change enhances the overall clarity of the code without altering functionality. --- src/openai/_utils/_streams.py | 1 + 1 file changed, 1 insertion(+) diff --git a/src/openai/_utils/_streams.py b/src/openai/_utils/_streams.py index 65408c3c9d..d5c2d55c6c 100644 --- a/src/openai/_utils/_streams.py +++ b/src/openai/_utils/_streams.py @@ -21,6 +21,7 @@ def drain_sync_iterator(iterator: Iterator[Any], timeout_ms: int = 50) -> None: server holds connection open after [DONE]. """ import time + deadline = time.monotonic() + (timeout_ms / 1000.0) try: while time.monotonic() < deadline: From f27ac4044fd403861f4a0ad90584d3ae9cfbb8e1 Mon Sep 17 00:00:00 2001 From: vrs-darkness Date: Fri, 14 Aug 2026 19:33:19 +0530 Subject: [PATCH 07/21] refactor(streaming): improve byte iterator handling in Stream classes - Introduced `_byte_iterator` attributes in both `Stream` and `AsyncStream` classes to store byte iterators, enhancing the efficiency of event iteration. - Updated `_iter_events` methods to initialize `_byte_iterator` only when necessary, optimizing resource usage. - Adjusted `__stream__` methods to utilize the new `_iter_events` methods for better clarity and performance. This change refines the handling of byte streams, ensuring more efficient iteration and resource management. --- src/openai/_streaming.py | 22 ++++++++++++++-------- 1 file changed, 14 insertions(+), 8 deletions(-) diff --git a/src/openai/_streaming.py b/src/openai/_streaming.py index 67e439333f..2179c6c324 100644 --- a/src/openai/_streaming.py +++ b/src/openai/_streaming.py @@ -40,6 +40,7 @@ def __init__( self._client = client self._options = options self._decoder = client._make_sse_decoder() + self._byte_iterator: Iterator[bytes] | None = None self._iterator = self.__stream__() def __next__(self) -> _T: @@ -50,14 +51,16 @@ def __iter__(self) -> Iterator[_T]: yield item def _iter_events(self) -> Iterator[ServerSentEvent]: - yield from self._decoder.iter_bytes(self.response.iter_bytes()) + if self._byte_iterator is None: + self._byte_iterator = self.response.iter_bytes() + yield from self._decoder.iter_bytes(self._byte_iterator) def __stream__(self) -> Iterator[_T]: cast_to = cast(Any, self._cast_to) response = self.response process_data = self._client._process_response_data - byte_iterator = response.iter_bytes() - iterator = self._decoder.iter_bytes(byte_iterator) + self._byte_iterator = response.iter_bytes() + iterator = self._iter_events() try: for sse in iterator: @@ -112,7 +115,7 @@ def __stream__(self) -> Iterator[_T]: # Drain raw bytes (not decoded events) so heartbeat comments don't cause unbounded waits. # Drain in finally so [DONE] doesn't block the caller. try: - drain_sync_iterator(byte_iterator, timeout_ms=50) + drain_sync_iterator(self._byte_iterator, timeout_ms=50) except (httpx2.HTTPError, UnicodeError): pass response.close() @@ -157,6 +160,7 @@ def __init__( self._client = client self._options = options self._decoder = client._make_sse_decoder() + self._byte_iterator: AsyncIterator[bytes] | None = None self._iterator = self.__stream__() async def __anext__(self) -> _T: @@ -167,15 +171,17 @@ async def __aiter__(self) -> AsyncIterator[_T]: yield item async def _iter_events(self) -> AsyncIterator[ServerSentEvent]: - async for sse in self._decoder.aiter_bytes(self.response.aiter_bytes()): + if self._byte_iterator is None: + self._byte_iterator = self.response.aiter_bytes() + async for sse in self._decoder.aiter_bytes(self._byte_iterator): yield sse async def __stream__(self) -> AsyncIterator[_T]: cast_to = cast(Any, self._cast_to) response = self.response process_data = self._client._process_response_data - byte_iterator = response.aiter_bytes() - iterator = self._decoder.aiter_bytes(byte_iterator) + self._byte_iterator = response.aiter_bytes() + iterator = self._iter_events() try: async for sse in iterator: @@ -230,7 +236,7 @@ async def __stream__(self) -> AsyncIterator[_T]: # Drain raw bytes (not decoded events) so heartbeat comments don't cause unbounded waits. # Drain in finally so [DONE] doesn't block the caller. try: - await drain_async_iterator(byte_iterator, timeout_ms=50) + await drain_async_iterator(self._byte_iterator, timeout_ms=50) except (httpx2.HTTPError, UnicodeError): pass await response.aclose() From 54493eb32947f49d39df75a62e6e375737b55a59 Mon Sep 17 00:00:00 2001 From: vrs-darkness Date: Fri, 14 Aug 2026 19:40:19 +0530 Subject: [PATCH 08/21] refactor(streaming): streamline error handling in iterator draining - Removed unnecessary try-except blocks in `drain_sync_iterator` and `drain_async_iterator` functions to simplify error handling. - Enhanced the logic to ensure that exceptions are handled more cleanly, improving code readability and maintainability. This change refines the iterator draining process, ensuring clearer error management while maintaining functionality. --- src/openai/_streaming.py | 10 ++------- src/openai/_utils/_streams.py | 38 ++++++++++++++++++----------------- 2 files changed, 22 insertions(+), 26 deletions(-) diff --git a/src/openai/_streaming.py b/src/openai/_streaming.py index 2179c6c324..90c6ceaf1c 100644 --- a/src/openai/_streaming.py +++ b/src/openai/_streaming.py @@ -114,10 +114,7 @@ def __stream__(self) -> Iterator[_T]: # [DONE] is already terminal for callers; drain failures must not fail the stream. # Drain raw bytes (not decoded events) so heartbeat comments don't cause unbounded waits. # Drain in finally so [DONE] doesn't block the caller. - try: - drain_sync_iterator(self._byte_iterator, timeout_ms=50) - except (httpx2.HTTPError, UnicodeError): - pass + drain_sync_iterator(self._byte_iterator, timeout_ms=50) response.close() def __enter__(self) -> Self: @@ -235,10 +232,7 @@ async def __stream__(self) -> AsyncIterator[_T]: # [DONE] is already terminal for callers; drain failures must not fail the stream. # Drain raw bytes (not decoded events) so heartbeat comments don't cause unbounded waits. # Drain in finally so [DONE] doesn't block the caller. - try: - await drain_async_iterator(self._byte_iterator, timeout_ms=50) - except (httpx2.HTTPError, UnicodeError): - pass + await drain_async_iterator(self._byte_iterator, timeout_ms=50) await response.aclose() async def __aenter__(self) -> Self: diff --git a/src/openai/_utils/_streams.py b/src/openai/_utils/_streams.py index d5c2d55c6c..605130edf4 100644 --- a/src/openai/_utils/_streams.py +++ b/src/openai/_utils/_streams.py @@ -23,14 +23,13 @@ def drain_sync_iterator(iterator: Iterator[Any], timeout_ms: int = 50) -> None: import time deadline = time.monotonic() + (timeout_ms / 1000.0) - try: - while time.monotonic() < deadline: - try: - next(iterator) - except StopIteration: - return - except Exception: - pass + while time.monotonic() < deadline: + try: + next(iterator) + except StopIteration: + return + except Exception: + break async def drain_async_iterator(iterator: AsyncIterator[Any], timeout_ms: int = 50) -> None: @@ -40,13 +39,16 @@ async def drain_async_iterator(iterator: AsyncIterator[Any], timeout_ms: int = 5 reuse, but gives up after timeout_ms to avoid indefinite blocking when server holds connection open after [DONE]. """ - try: - while True: - try: - await asyncio.wait_for(iterator.__anext__(), timeout=timeout_ms / 1000.0) - except StopAsyncIteration: - return - except asyncio.TimeoutError: - pass - except Exception: - pass + deadline = asyncio.get_event_loop().time() + (timeout_ms / 1000.0) + while True: + remaining = deadline - asyncio.get_event_loop().time() + if remaining <= 0: + break + try: + await asyncio.wait_for(iterator.__anext__(), timeout=remaining) + except StopAsyncIteration: + return + except asyncio.TimeoutError: + break + except Exception: + break From cf37030f6a6ff8f17c3cf95346fb3891a5ec2070 Mon Sep 17 00:00:00 2001 From: vrs-darkness Date: Fri, 14 Aug 2026 19:43:44 +0530 Subject: [PATCH 09/21] refactor(streaming): update comments for iterator draining - Simplified comments in `Stream` and `AsyncStream` classes to clarify the purpose of draining remaining bytes for connection reuse. - Adjusted the timeout handling in `drain_async_iterator` to use the current running event loop, enhancing clarity and consistency. This change improves the readability of the code while maintaining the functionality of the iterator draining process. --- src/openai/_streaming.py | 10 ++-------- src/openai/_utils/_streams.py | 5 +++-- 2 files changed, 5 insertions(+), 10 deletions(-) diff --git a/src/openai/_streaming.py b/src/openai/_streaming.py index 90c6ceaf1c..87e585a9a6 100644 --- a/src/openai/_streaming.py +++ b/src/openai/_streaming.py @@ -110,10 +110,7 @@ def __stream__(self) -> Iterator[_T]: response=response, ) finally: - # Best-effort timeout-bounded drain of raw bytes to enable connection reuse. - # [DONE] is already terminal for callers; drain failures must not fail the stream. - # Drain raw bytes (not decoded events) so heartbeat comments don't cause unbounded waits. - # Drain in finally so [DONE] doesn't block the caller. + # Drain remaining body to enable connection reuse; best-effort with 50ms timeout. drain_sync_iterator(self._byte_iterator, timeout_ms=50) response.close() @@ -228,10 +225,7 @@ async def __stream__(self) -> AsyncIterator[_T]: response=response, ) finally: - # Best-effort timeout-bounded drain of raw bytes to enable connection reuse. - # [DONE] is already terminal for callers; drain failures must not fail the stream. - # Drain raw bytes (not decoded events) so heartbeat comments don't cause unbounded waits. - # Drain in finally so [DONE] doesn't block the caller. + # Drain remaining body to enable connection reuse; best-effort with 50ms timeout. await drain_async_iterator(self._byte_iterator, timeout_ms=50) await response.aclose() diff --git a/src/openai/_utils/_streams.py b/src/openai/_utils/_streams.py index 605130edf4..2cd3aea0c5 100644 --- a/src/openai/_utils/_streams.py +++ b/src/openai/_utils/_streams.py @@ -39,9 +39,10 @@ async def drain_async_iterator(iterator: AsyncIterator[Any], timeout_ms: int = 5 reuse, but gives up after timeout_ms to avoid indefinite blocking when server holds connection open after [DONE]. """ - deadline = asyncio.get_event_loop().time() + (timeout_ms / 1000.0) + loop = asyncio.get_running_loop() + deadline = loop.time() + (timeout_ms / 1000.0) while True: - remaining = deadline - asyncio.get_event_loop().time() + remaining = deadline - loop.time() if remaining <= 0: break try: From b94d8a3a2a74a9eb711351b55c5fc73d9cf72638 Mon Sep 17 00:00:00 2001 From: vrs-darkness Date: Fri, 14 Aug 2026 19:46:28 +0530 Subject: [PATCH 10/21] refactor(streaming): enhance iterator draining with optional handling - Updated `drain_sync_iterator` and `drain_async_iterator` functions to accept optional iterators, preventing indefinite blocking when a None value is passed. - Modified test cases to reflect changes in iterator behavior, ensuring that slow iterators trigger timeouts as expected. This change improves the robustness of iterator draining by handling None values gracefully and ensuring proper timeout behavior during draining. --- src/openai/_utils/_streams.py | 13 ++++++++---- tests/test_streaming.py | 37 +++++++++++++++-------------------- 2 files changed, 25 insertions(+), 25 deletions(-) diff --git a/src/openai/_utils/_streams.py b/src/openai/_utils/_streams.py index 2cd3aea0c5..e9dd493137 100644 --- a/src/openai/_utils/_streams.py +++ b/src/openai/_utils/_streams.py @@ -1,5 +1,6 @@ +import time import asyncio -from typing import Any +from typing import Any, Optional from typing_extensions import Iterator, AsyncIterator @@ -13,14 +14,15 @@ async def consume_async_iterator(iterator: AsyncIterator[Any]) -> None: ... -def drain_sync_iterator(iterator: Iterator[Any], timeout_ms: int = 50) -> None: +def drain_sync_iterator(iterator: Optional[Iterator[Any]], timeout_ms: int = 50) -> None: """Drain trailing bytes from iterator with bounded timeout. Attempts to drain all remaining items from iterator to enable connection reuse, but gives up after timeout_ms to avoid indefinite blocking when server holds connection open after [DONE]. """ - import time + if iterator is None: + return deadline = time.monotonic() + (timeout_ms / 1000.0) while time.monotonic() < deadline: @@ -32,13 +34,16 @@ def drain_sync_iterator(iterator: Iterator[Any], timeout_ms: int = 50) -> None: break -async def drain_async_iterator(iterator: AsyncIterator[Any], timeout_ms: int = 50) -> None: +async def drain_async_iterator(iterator: Optional[AsyncIterator[Any]], timeout_ms: int = 50) -> None: """Drain trailing bytes from async iterator with bounded timeout. Attempts to drain all remaining items from iterator to enable connection reuse, but gives up after timeout_ms to avoid indefinite blocking when server holds connection open after [DONE]. """ + if iterator is None: + return + loop = asyncio.get_running_loop() deadline = loop.time() + (timeout_ms / 1000.0) while True: diff --git a/tests/test_streaming.py b/tests/test_streaming.py index 30bf4ad4d0..d85a6dd4b7 100644 --- a/tests/test_streaming.py +++ b/tests/test_streaming.py @@ -297,15 +297,14 @@ def body() -> Iterator[bytes]: async def test_drain_is_bounded_and_doesnt_block_indefinitely( sync: bool, client: OpenAI, async_client: AsyncOpenAI ) -> None: - """Drain after [DONE] must be bounded; trailing iterator that never closes must not block completion.""" + """Drain after [DONE] must timeout rather than wait for slow/infinite iterators.""" - class SyncInfiniteIterator: - """Sync iterator that yields indefinitely after content; simulates unbounded heartbeats.""" + class SyncSlowIterator: + """Sync iterator that yields slowly after content to force timeout.""" def __init__(self, content: Iterator[bytes]) -> None: self._content = content self._content_exhausted = False - self._yield_count = 0 def __iter__(self) -> Iterator[bytes]: return self @@ -316,19 +315,17 @@ def __next__(self) -> bytes: return next(self._content) except StopIteration: self._content_exhausted = True - # After content is exhausted, yield indefinitely (simulating held-open connection) - if self._yield_count > 1000: # Safety limit to prevent infinite test loops - raise StopIteration - self._yield_count += 1 + # After content is exhausted, yield slowly to force timeout during drain + import time + time.sleep(0.1) # 100ms per item; 50ms drain timeout will hit after <1 item return b": heartbeat\n\n" - class AsyncInfiniteIterator: - """Async iterator that yields indefinitely after content; simulates unbounded heartbeats.""" + class AsyncSlowIterator: + """Async iterator that yields slowly after content to force timeout.""" def __init__(self, content: AsyncIterator[bytes]) -> None: self._content = content self._content_exhausted = False - self._yield_count = 0 def __aiter__(self) -> AsyncIterator[bytes]: return self @@ -339,30 +336,28 @@ async def __anext__(self) -> bytes: return await self._content.__anext__() except StopAsyncIteration: self._content_exhausted = True - # After content is exhausted, yield indefinitely (simulating held-open connection) - if self._yield_count > 1000: # Safety limit to prevent infinite test loops - raise StopAsyncIteration - self._yield_count += 1 + # After content is exhausted, yield slowly to force timeout during drain + import asyncio + await asyncio.sleep(0.1) # 100ms per item; 50ms drain timeout will hit after <1 item return b": heartbeat\n\n" def body() -> Iterator[bytes]: yield b'data: {"foo":true}\n\n' yield b"data: [DONE]\n\n" - # More events after [DONE] that should not block completion if sync: - infinite_iter = SyncInfiniteIterator(body()) - response = httpx2.Response(200, content=infinite_iter) + slow_iter = SyncSlowIterator(body()) + response = httpx2.Response(200, content=slow_iter) stream: Stream[object] | AsyncStream[object] = Stream(cast_to=object, client=client, response=response) chunks = list(stream) else: async_body = to_aiter(body()) - infinite_iter = AsyncInfiniteIterator(async_body) - response = httpx2.Response(200, content=infinite_iter) + slow_iter = AsyncSlowIterator(async_body) + response = httpx2.Response(200, content=slow_iter) stream = AsyncStream(cast_to=object, client=async_client, response=response) chunks = [chunk async for chunk in stream] - # Stream should complete with just the first item, not block waiting for the infinite iterator + # Stream should complete quickly with just the first item, not wait for the slow iterator assert chunks == [{"foo": True}] assert response.is_closed is True From 34248fa53f316310830c494be2f0e8e979f0d3cd Mon Sep 17 00:00:00 2001 From: vrs-darkness Date: Fri, 14 Aug 2026 19:53:13 +0530 Subject: [PATCH 11/21] refactor(streaming): enhance synchronous iterator draining with threading - Modified `drain_sync_iterator` to run in a background thread, preventing blocking on `iterator.__next__()`. - Improved error handling within the draining process to ensure graceful termination on exceptions. This change optimizes the synchronous iterator draining process, enhancing responsiveness and robustness during iteration. --- src/openai/_utils/_streams.py | 22 ++++++++++++++++------ 1 file changed, 16 insertions(+), 6 deletions(-) diff --git a/src/openai/_utils/_streams.py b/src/openai/_utils/_streams.py index e9dd493137..df6384424b 100644 --- a/src/openai/_utils/_streams.py +++ b/src/openai/_utils/_streams.py @@ -1,5 +1,6 @@ import time import asyncio +import threading from typing import Any, Optional from typing_extensions import Iterator, AsyncIterator @@ -20,18 +21,27 @@ def drain_sync_iterator(iterator: Optional[Iterator[Any]], timeout_ms: int = 50) Attempts to drain all remaining items from iterator to enable connection reuse, but gives up after timeout_ms to avoid indefinite blocking when server holds connection open after [DONE]. + + Runs in background thread to prevent blocking on iterator.__next__(). """ if iterator is None: return - deadline = time.monotonic() + (timeout_ms / 1000.0) - while time.monotonic() < deadline: + def _drain() -> None: try: - next(iterator) - except StopIteration: - return + while True: + try: + next(iterator) + except StopIteration: + return + except Exception: + break except Exception: - break + pass + + thread = threading.Thread(target=_drain, daemon=True) + thread.start() + thread.join(timeout=timeout_ms / 1000.0) async def drain_async_iterator(iterator: Optional[AsyncIterator[Any]], timeout_ms: int = 50) -> None: From 4e25c1e015b4b8ccd01226ee65af57944ca789c2 Mon Sep 17 00:00:00 2001 From: vrs-darkness Date: Fri, 14 Aug 2026 19:53:18 +0530 Subject: [PATCH 12/21] refactor(streaming): remove unused import in _streams.py - Eliminated the unused `time` import from `_streams.py` to enhance code cleanliness and maintainability. This change contributes to a more streamlined codebase by removing unnecessary dependencies. --- src/openai/_utils/_streams.py | 1 - 1 file changed, 1 deletion(-) diff --git a/src/openai/_utils/_streams.py b/src/openai/_utils/_streams.py index df6384424b..98b321c561 100644 --- a/src/openai/_utils/_streams.py +++ b/src/openai/_utils/_streams.py @@ -1,4 +1,3 @@ -import time import asyncio import threading from typing import Any, Optional From c1dc31b190487dae78c30923850d701b0cf7fa0c Mon Sep 17 00:00:00 2001 From: vrs-darkness Date: Fri, 14 Aug 2026 19:57:46 +0530 Subject: [PATCH 13/21] refactor(streaming): integrate anyio for async iterator draining - Replaced asyncio-based timeout handling in `drain_async_iterator` with anyio's `move_on_after` for improved compatibility with both asyncio and Trio. - Streamlined the draining logic to utilize async iteration directly, enhancing code clarity and efficiency. This change enhances the flexibility of the async iterator draining process, ensuring better performance across different async backends. --- src/openai/_utils/_streams.py | 21 ++++++++------------- 1 file changed, 8 insertions(+), 13 deletions(-) diff --git a/src/openai/_utils/_streams.py b/src/openai/_utils/_streams.py index 98b321c561..04998c5f29 100644 --- a/src/openai/_utils/_streams.py +++ b/src/openai/_utils/_streams.py @@ -1,8 +1,9 @@ -import asyncio import threading from typing import Any, Optional from typing_extensions import Iterator, AsyncIterator +import anyio + def consume_sync_iterator(iterator: Iterator[Any]) -> None: for _ in iterator: @@ -49,21 +50,15 @@ async def drain_async_iterator(iterator: Optional[AsyncIterator[Any]], timeout_m Attempts to drain all remaining items from iterator to enable connection reuse, but gives up after timeout_ms to avoid indefinite blocking when server holds connection open after [DONE]. + + Uses anyio for async backend compatibility (works with asyncio and Trio). """ if iterator is None: return - loop = asyncio.get_running_loop() - deadline = loop.time() + (timeout_ms / 1000.0) - while True: - remaining = deadline - loop.time() - if remaining <= 0: - break + with anyio.move_on_after(timeout_ms / 1000.0): try: - await asyncio.wait_for(iterator.__anext__(), timeout=remaining) - except StopAsyncIteration: - return - except asyncio.TimeoutError: - break + async for _ in iterator: + pass except Exception: - break + pass From d9b541170c105714fed8943f85f14d7bf72602c2 Mon Sep 17 00:00:00 2001 From: vrs-darkness Date: Fri, 14 Aug 2026 20:12:12 +0530 Subject: [PATCH 14/21] refactor(streaming): enhance iterator draining with response handling - Updated `drain_sync_iterator` and `drain_async_iterator` to accept a `response` parameter, allowing for graceful closure of responses if draining times out. - Improved error handling during response closure to prevent unhandled exceptions, ensuring robust behavior in both synchronous and asynchronous contexts. This change enhances the reliability of iterator draining by ensuring that blocked reads are properly interrupted, improving overall stream management. --- src/openai/_streaming.py | 16 ++++++++++++---- src/openai/_utils/_streams.py | 35 ++++++++++++++++++++++++++++------- tests/test_streaming.py | 2 ++ 3 files changed, 42 insertions(+), 11 deletions(-) diff --git a/src/openai/_streaming.py b/src/openai/_streaming.py index 87e585a9a6..90337bb702 100644 --- a/src/openai/_streaming.py +++ b/src/openai/_streaming.py @@ -111,8 +111,12 @@ def __stream__(self) -> Iterator[_T]: ) finally: # Drain remaining body to enable connection reuse; best-effort with 50ms timeout. - drain_sync_iterator(self._byte_iterator, timeout_ms=50) - response.close() + # If drain times out, response is closed to interrupt any blocked reads. + drain_sync_iterator(self._byte_iterator, response=response, timeout_ms=50) + try: + response.close() + except Exception: + pass def __enter__(self) -> Self: return self @@ -226,8 +230,12 @@ async def __stream__(self) -> AsyncIterator[_T]: ) finally: # Drain remaining body to enable connection reuse; best-effort with 50ms timeout. - await drain_async_iterator(self._byte_iterator, timeout_ms=50) - await response.aclose() + # If drain times out, response is closed to interrupt any blocked reads. + await drain_async_iterator(self._byte_iterator, response=response, timeout_ms=50) + try: + await response.aclose() + except Exception: + pass async def __aenter__(self) -> Self: return self diff --git a/src/openai/_utils/_streams.py b/src/openai/_utils/_streams.py index 04998c5f29..c218c0c111 100644 --- a/src/openai/_utils/_streams.py +++ b/src/openai/_utils/_streams.py @@ -15,7 +15,9 @@ async def consume_async_iterator(iterator: AsyncIterator[Any]) -> None: ... -def drain_sync_iterator(iterator: Optional[Iterator[Any]], timeout_ms: int = 50) -> None: +def drain_sync_iterator( + iterator: Optional[Iterator[Any]], response: Any = None, timeout_ms: int = 50 +) -> None: """Drain trailing bytes from iterator with bounded timeout. Attempts to drain all remaining items from iterator to enable connection @@ -23,6 +25,7 @@ def drain_sync_iterator(iterator: Optional[Iterator[Any]], timeout_ms: int = 50) server holds connection open after [DONE]. Runs in background thread to prevent blocking on iterator.__next__(). + If drain times out, closes response from main thread to interrupt blocked reads. """ if iterator is None: return @@ -43,8 +46,17 @@ def _drain() -> None: thread.start() thread.join(timeout=timeout_ms / 1000.0) + # If thread still alive after timeout, close response to interrupt blocked read + if thread.is_alive() and response is not None: + try: + response.close() + except Exception: + pass + -async def drain_async_iterator(iterator: Optional[AsyncIterator[Any]], timeout_ms: int = 50) -> None: +async def drain_async_iterator( + iterator: Optional[AsyncIterator[Any]], response: Any = None, timeout_ms: int = 50 +) -> None: """Drain trailing bytes from async iterator with bounded timeout. Attempts to drain all remaining items from iterator to enable connection @@ -52,13 +64,22 @@ async def drain_async_iterator(iterator: Optional[AsyncIterator[Any]], timeout_m server holds connection open after [DONE]. Uses anyio for async backend compatibility (works with asyncio and Trio). + If drain times out, closes response to interrupt blocked reads. """ if iterator is None: return - with anyio.move_on_after(timeout_ms / 1000.0): - try: - async for _ in iterator: + try: + with anyio.move_on_after(timeout_ms / 1000.0): + try: + async for _ in iterator: + pass + except Exception: + pass + finally: + # Close response to interrupt any blocked iterator reads on timeout + if response is not None: + try: + await response.aclose() + except Exception: pass - except Exception: - pass diff --git a/tests/test_streaming.py b/tests/test_streaming.py index d85a6dd4b7..cb3041ee12 100644 --- a/tests/test_streaming.py +++ b/tests/test_streaming.py @@ -317,6 +317,7 @@ def __next__(self) -> bytes: self._content_exhausted = True # After content is exhausted, yield slowly to force timeout during drain import time + time.sleep(0.1) # 100ms per item; 50ms drain timeout will hit after <1 item return b": heartbeat\n\n" @@ -338,6 +339,7 @@ async def __anext__(self) -> bytes: self._content_exhausted = True # After content is exhausted, yield slowly to force timeout during drain import asyncio + await asyncio.sleep(0.1) # 100ms per item; 50ms drain timeout will hit after <1 item return b": heartbeat\n\n" From ea219c4d61d2c2d2187108dd8994c80f5d48b05a Mon Sep 17 00:00:00 2001 From: vrs-darkness Date: Fri, 14 Aug 2026 20:17:09 +0530 Subject: [PATCH 15/21] refactor(streaming): improve handling of iterator draining post [DONE] event - Introduced a `_done_seen` attribute in both `Stream` and `AsyncStream` classes to track the completion of the stream. - Updated the draining logic to only execute if the [DONE] event has been seen, preventing unnecessary consumption of data on early exits. - Enhanced comments to clarify the purpose of the changes in the draining process. This change optimizes resource management during stream termination, ensuring that abandoned data is not consumed unnecessarily. --- src/openai/_streaming.py | 18 ++++++++++++------ 1 file changed, 12 insertions(+), 6 deletions(-) diff --git a/src/openai/_streaming.py b/src/openai/_streaming.py index 90337bb702..1b2e2f22ee 100644 --- a/src/openai/_streaming.py +++ b/src/openai/_streaming.py @@ -41,6 +41,7 @@ def __init__( self._options = options self._decoder = client._make_sse_decoder() self._byte_iterator: Iterator[bytes] | None = None + self._done_seen = False self._iterator = self.__stream__() def __next__(self) -> _T: @@ -65,6 +66,7 @@ def __stream__(self) -> Iterator[_T]: try: for sse in iterator: if sse.data.startswith("[DONE]"): + self._done_seen = True break # we have to special case the Assistants `thread.` events since we won't have an "event" key in the data @@ -110,9 +112,10 @@ def __stream__(self) -> Iterator[_T]: response=response, ) finally: - # Drain remaining body to enable connection reuse; best-effort with 50ms timeout. - # If drain times out, response is closed to interrupt any blocked reads. - drain_sync_iterator(self._byte_iterator, response=response, timeout_ms=50) + # Drain remaining body only after [DONE] to enable connection reuse. + # On early exit (before [DONE]), skip drain to avoid consuming abandoned data. + if self._done_seen: + drain_sync_iterator(self._byte_iterator, response=response, timeout_ms=50) try: response.close() except Exception: @@ -159,6 +162,7 @@ def __init__( self._options = options self._decoder = client._make_sse_decoder() self._byte_iterator: AsyncIterator[bytes] | None = None + self._done_seen = False self._iterator = self.__stream__() async def __anext__(self) -> _T: @@ -184,6 +188,7 @@ async def __stream__(self) -> AsyncIterator[_T]: try: async for sse in iterator: if sse.data.startswith("[DONE]"): + self._done_seen = True break # we have to special case the Assistants `thread.` events since we won't have an "event" key in the data @@ -229,9 +234,10 @@ async def __stream__(self) -> AsyncIterator[_T]: response=response, ) finally: - # Drain remaining body to enable connection reuse; best-effort with 50ms timeout. - # If drain times out, response is closed to interrupt any blocked reads. - await drain_async_iterator(self._byte_iterator, response=response, timeout_ms=50) + # Drain remaining body only after [DONE] to enable connection reuse. + # On early exit (before [DONE]), skip drain to avoid consuming abandoned data. + if self._done_seen: + await drain_async_iterator(self._byte_iterator, response=response, timeout_ms=50) try: await response.aclose() except Exception: From 10adc436a42dbbd24c7f14d838e77b963e3d08cc Mon Sep 17 00:00:00 2001 From: vrs-darkness Date: Fri, 14 Aug 2026 20:17:40 +0530 Subject: [PATCH 16/21] refactor(streaming): streamline function signatures for iterator draining - Removed unnecessary line breaks in the function signatures of `drain_sync_iterator` and `drain_async_iterator` to improve code readability. - This change enhances the clarity of the code without altering the functionality of the iterator draining process. --- src/openai/_utils/_streams.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/src/openai/_utils/_streams.py b/src/openai/_utils/_streams.py index c218c0c111..dd19ad04c7 100644 --- a/src/openai/_utils/_streams.py +++ b/src/openai/_utils/_streams.py @@ -15,9 +15,7 @@ async def consume_async_iterator(iterator: AsyncIterator[Any]) -> None: ... -def drain_sync_iterator( - iterator: Optional[Iterator[Any]], response: Any = None, timeout_ms: int = 50 -) -> None: +def drain_sync_iterator(iterator: Optional[Iterator[Any]], response: Any = None, timeout_ms: int = 50) -> None: """Drain trailing bytes from iterator with bounded timeout. Attempts to drain all remaining items from iterator to enable connection From 5055298fd5b571fd6b192c132f3cf94ed07b3730 Mon Sep 17 00:00:00 2001 From: vrs-darkness Date: Fri, 14 Aug 2026 20:21:31 +0530 Subject: [PATCH 17/21] refactor(streaming): improve thread handling in synchronous iterator draining - Added a join operation for the background thread in `drain_sync_iterator` to ensure it exits gracefully after an interrupt. - This change enhances the robustness of the synchronous iterator draining process by providing adequate time for thread termination, improving overall resource management. --- src/openai/_utils/_streams.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/src/openai/_utils/_streams.py b/src/openai/_utils/_streams.py index dd19ad04c7..099f7002a6 100644 --- a/src/openai/_utils/_streams.py +++ b/src/openai/_utils/_streams.py @@ -50,6 +50,8 @@ def _drain() -> None: response.close() except Exception: pass + # Rejoin to give thread time to exit after interrupt + thread.join(timeout=timeout_ms / 1000.0) async def drain_async_iterator( From 58d222d89bcb3737dc18af0e610dc120d99ed88b Mon Sep 17 00:00:00 2001 From: vrs-darkness Date: Fri, 14 Aug 2026 20:51:47 +0530 Subject: [PATCH 18/21] refactor(streaming): simplify iterator draining logic and enhance response handling - Removed the `_done_seen` attribute from `Stream` and `AsyncStream` classes, streamlining the logic for handling the [DONE] event during iteration. - Updated the draining process to always close the response after the stream ends, ensuring proper resource management and connection reuse. - Adjusted comments for clarity regarding the purpose of draining remaining bytes. This change improves the efficiency and reliability of the streaming process by simplifying the code and enhancing response closure behavior. --- src/openai/_streaming.py | 21 +++------------ src/openai/_utils/__init__.py | 1 - src/openai/_utils/_streams.py | 49 +++-------------------------------- tests/test_streaming.py | 33 +++++++++++++++-------- 4 files changed, 29 insertions(+), 75 deletions(-) diff --git a/src/openai/_streaming.py b/src/openai/_streaming.py index 1b2e2f22ee..e9e94baabe 100644 --- a/src/openai/_streaming.py +++ b/src/openai/_streaming.py @@ -9,7 +9,7 @@ import httpx2 -from ._utils import is_mapping, drain_sync_iterator, drain_async_iterator, extract_type_var_from_base +from ._utils import is_mapping, drain_async_iterator, extract_type_var_from_base from ._exceptions import APIError if TYPE_CHECKING: @@ -41,7 +41,6 @@ def __init__( self._options = options self._decoder = client._make_sse_decoder() self._byte_iterator: Iterator[bytes] | None = None - self._done_seen = False self._iterator = self.__stream__() def __next__(self) -> _T: @@ -66,7 +65,6 @@ def __stream__(self) -> Iterator[_T]: try: for sse in iterator: if sse.data.startswith("[DONE]"): - self._done_seen = True break # we have to special case the Assistants `thread.` events since we won't have an "event" key in the data @@ -112,14 +110,7 @@ def __stream__(self) -> Iterator[_T]: response=response, ) finally: - # Drain remaining body only after [DONE] to enable connection reuse. - # On early exit (before [DONE]), skip drain to avoid consuming abandoned data. - if self._done_seen: - drain_sync_iterator(self._byte_iterator, response=response, timeout_ms=50) - try: - response.close() - except Exception: - pass + response.close() def __enter__(self) -> Self: return self @@ -162,7 +153,6 @@ def __init__( self._options = options self._decoder = client._make_sse_decoder() self._byte_iterator: AsyncIterator[bytes] | None = None - self._done_seen = False self._iterator = self.__stream__() async def __anext__(self) -> _T: @@ -188,7 +178,6 @@ async def __stream__(self) -> AsyncIterator[_T]: try: async for sse in iterator: if sse.data.startswith("[DONE]"): - self._done_seen = True break # we have to special case the Assistants `thread.` events since we won't have an "event" key in the data @@ -234,10 +223,8 @@ async def __stream__(self) -> AsyncIterator[_T]: response=response, ) finally: - # Drain remaining body only after [DONE] to enable connection reuse. - # On early exit (before [DONE]), skip drain to avoid consuming abandoned data. - if self._done_seen: - await drain_async_iterator(self._byte_iterator, response=response, timeout_ms=50) + # Bounds cancellation-cooperative async drain for connection reuse. + await drain_async_iterator(self._byte_iterator, response=response, timeout_ms=50) try: await response.aclose() except Exception: diff --git a/src/openai/_utils/__init__.py b/src/openai/_utils/__init__.py index a91fc2f6c1..8951d8e502 100644 --- a/src/openai/_utils/__init__.py +++ b/src/openai/_utils/__init__.py @@ -53,7 +53,6 @@ extract_type_var_from_base as extract_type_var_from_base, ) from ._streams import ( - drain_sync_iterator as drain_sync_iterator, drain_async_iterator as drain_async_iterator, consume_sync_iterator as consume_sync_iterator, consume_async_iterator as consume_async_iterator, diff --git a/src/openai/_utils/_streams.py b/src/openai/_utils/_streams.py index 099f7002a6..28119ea558 100644 --- a/src/openai/_utils/_streams.py +++ b/src/openai/_utils/_streams.py @@ -1,4 +1,3 @@ -import threading from typing import Any, Optional from typing_extensions import Iterator, AsyncIterator @@ -15,56 +14,14 @@ async def consume_async_iterator(iterator: AsyncIterator[Any]) -> None: ... -def drain_sync_iterator(iterator: Optional[Iterator[Any]], response: Any = None, timeout_ms: int = 50) -> None: - """Drain trailing bytes from iterator with bounded timeout. - - Attempts to drain all remaining items from iterator to enable connection - reuse, but gives up after timeout_ms to avoid indefinite blocking when - server holds connection open after [DONE]. - - Runs in background thread to prevent blocking on iterator.__next__(). - If drain times out, closes response from main thread to interrupt blocked reads. - """ - if iterator is None: - return - - def _drain() -> None: - try: - while True: - try: - next(iterator) - except StopIteration: - return - except Exception: - break - except Exception: - pass - - thread = threading.Thread(target=_drain, daemon=True) - thread.start() - thread.join(timeout=timeout_ms / 1000.0) - - # If thread still alive after timeout, close response to interrupt blocked read - if thread.is_alive() and response is not None: - try: - response.close() - except Exception: - pass - # Rejoin to give thread time to exit after interrupt - thread.join(timeout=timeout_ms / 1000.0) - - async def drain_async_iterator( iterator: Optional[AsyncIterator[Any]], response: Any = None, timeout_ms: int = 50 ) -> None: """Drain trailing bytes from async iterator with bounded timeout. - Attempts to drain all remaining items from iterator to enable connection - reuse, but gives up after timeout_ms to avoid indefinite blocking when - server holds connection open after [DONE]. - - Uses anyio for async backend compatibility (works with asyncio and Trio). - If drain times out, closes response to interrupt blocked reads. + Bounds cancellation-cooperative async streams to enable connection reuse after [DONE]. + Uses anyio for backend compatibility (asyncio and Trio). Cannot guarantee termination + for custom iterators that block the event loop or suppress cancellation. """ if iterator is None: return diff --git a/tests/test_streaming.py b/tests/test_streaming.py index cb3041ee12..3e1fc23f86 100644 --- a/tests/test_streaming.py +++ b/tests/test_streaming.py @@ -216,10 +216,25 @@ def body() -> Iterator[bytes]: assert sse.json() == {"content": "известни"} +def test_done_closes_response_sync(client: OpenAI) -> None: + """Sync stream closes response after [DONE].""" + + def body() -> Iterator[bytes]: + yield b'data: {"foo":true}\n\n' + yield b"data: [DONE]\n\n" + yield b": trailing comment after done\n\n" + + response = httpx2.Response(200, content=body()) + stream: Stream[object] | AsyncStream[object] = Stream(cast_to=object, client=client, response=response) + chunks = list(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_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.""" +async def test_done_drains_remaining_body_async(async_client: AsyncOpenAI) -> None: + """After [DONE], async drains remaining body for connection reuse.""" exhausted = False def body() -> Iterator[bytes]: @@ -229,14 +244,10 @@ def body() -> Iterator[bytes]: yield b": trailing comment after done\n\n" exhausted = True - response = httpx2.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] + async_body = to_aiter(body()) + response = httpx2.Response(200, content=async_body) + 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 From 990ed6fb81215df2ea4aeba0aa2e346c34db62f1 Mon Sep 17 00:00:00 2001 From: vrs-darkness Date: Fri, 14 Aug 2026 21:01:59 +0530 Subject: [PATCH 19/21] refactor(streaming): simplify byte iterator handling in Stream class - Removed the `_byte_iterator` attribute from the `Stream` class, streamlining the iteration logic. - Updated the `_iter_events` method to directly yield from the response's byte iterator, enhancing code clarity. - Added a comment to ensure the response is closed properly, improving resource management. This change simplifies the streaming process by reducing unnecessary state management and ensuring proper closure of resources. --- src/openai/_streaming.py | 7 ++----- 1 file changed, 2 insertions(+), 5 deletions(-) diff --git a/src/openai/_streaming.py b/src/openai/_streaming.py index e9e94baabe..1898da0481 100644 --- a/src/openai/_streaming.py +++ b/src/openai/_streaming.py @@ -40,7 +40,6 @@ def __init__( self._client = client self._options = options self._decoder = client._make_sse_decoder() - self._byte_iterator: Iterator[bytes] | None = None self._iterator = self.__stream__() def __next__(self) -> _T: @@ -51,15 +50,12 @@ def __iter__(self) -> Iterator[_T]: yield item def _iter_events(self) -> Iterator[ServerSentEvent]: - if self._byte_iterator is None: - self._byte_iterator = self.response.iter_bytes() - yield from self._decoder.iter_bytes(self._byte_iterator) + yield from self._decoder.iter_bytes(self.response.iter_bytes()) def __stream__(self) -> Iterator[_T]: cast_to = cast(Any, self._cast_to) response = self.response process_data = self._client._process_response_data - self._byte_iterator = response.iter_bytes() iterator = self._iter_events() try: @@ -110,6 +106,7 @@ def __stream__(self) -> Iterator[_T]: response=response, ) finally: + # Ensure the response is closed even if the consumer doesn't read all data response.close() def __enter__(self) -> Self: From 33fdf50406f035975e2892b74550252a9bac241a Mon Sep 17 00:00:00 2001 From: vrs-darkness Date: Fri, 14 Aug 2026 21:11:43 +0530 Subject: [PATCH 20/21] fix(streaming): ensure proper draining behavior on early exit from stream - Introduced a `_done_seen` attribute in the `AsyncStream` class to track whether the [DONE] event has been encountered during streaming. - Updated the draining logic to only execute if `_done_seen` is true, preventing unnecessary consumption of trailing data when exiting early. - Added a new test to verify that early exits without seeing [DONE] do not trigger the drain function, ensuring efficient resource management. This change enhances the reliability of the streaming process by ensuring that only relevant data is consumed during termination. --- src/openai/_streaming.py | 5 ++++- tests/test_streaming.py | 47 ++++++++++++++++++++++++++++++++++++++++ 2 files changed, 51 insertions(+), 1 deletion(-) diff --git a/src/openai/_streaming.py b/src/openai/_streaming.py index 1898da0481..2e93322aa8 100644 --- a/src/openai/_streaming.py +++ b/src/openai/_streaming.py @@ -150,6 +150,7 @@ def __init__( self._options = options self._decoder = client._make_sse_decoder() self._byte_iterator: AsyncIterator[bytes] | None = None + self._done_seen = False self._iterator = self.__stream__() async def __anext__(self) -> _T: @@ -175,6 +176,7 @@ async def __stream__(self) -> AsyncIterator[_T]: try: async for sse in iterator: if sse.data.startswith("[DONE]"): + self._done_seen = True break # we have to special case the Assistants `thread.` events since we won't have an "event" key in the data @@ -221,7 +223,8 @@ async def __stream__(self) -> AsyncIterator[_T]: ) finally: # Bounds cancellation-cooperative async drain for connection reuse. - await drain_async_iterator(self._byte_iterator, response=response, timeout_ms=50) + if self._done_seen: + await drain_async_iterator(self._byte_iterator, response=response, timeout_ms=50) try: await response.aclose() except Exception: diff --git a/tests/test_streaming.py b/tests/test_streaming.py index 3e1fc23f86..45691b58ce 100644 --- a/tests/test_streaming.py +++ b/tests/test_streaming.py @@ -1,6 +1,7 @@ from __future__ import annotations from collections.abc import Iterator, AsyncIterator +from typing import Any import httpx2 import pytest @@ -254,6 +255,52 @@ def body() -> Iterator[bytes]: assert response.is_closed is True +@pytest.mark.asyncio +async def test_early_exit_without_done_doesnt_drain_async(async_client: AsyncOpenAI) -> None: + """Early exit before [DONE] should not consume trailing body and close promptly.""" + drained = False + exhausted = False + + async def patched_drain(*_args: Any, **_kwargs: Any) -> None: + """Track if drain was called.""" + nonlocal drained + drained = True + + def body() -> Iterator[bytes]: + nonlocal exhausted + yield b'data: {"foo":true}\n\n' + # No [DONE] sent, consumer will exit early + yield b": trailing comment that should not be consumed\n\n" + exhausted = True + + async_body = to_aiter(body()) + response = httpx2.Response(200, content=async_body) + + # Patch drain_async_iterator to track if it's called + import openai._streaming as streaming_module + original_drain = streaming_module.drain_async_iterator + streaming_module.drain_async_iterator = patched_drain + + try: + stream = AsyncStream(cast_to=object, client=async_client, response=response) + + # Only consume the first chunk, exit early before [DONE] + first_chunk = await stream.__anext__() + assert first_chunk == {"foo": True} + + # Trailing body should NOT be exhausted since _done_seen is False + assert exhausted is False + + # Delete stream to finalize the generator and trigger finally block + del stream + + # Drain should NOT have been called since we didn't see [DONE] + assert drained is False + finally: + # Restore original drain function + streaming_module.drain_async_iterator = original_drain + + @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: From cab1db43e78f14f398f6934f9f67f0ee1432aded Mon Sep 17 00:00:00 2001 From: vrs-darkness Date: Fri, 14 Aug 2026 21:30:29 +0530 Subject: [PATCH 21/21] fix(streaming): guard async drain with _done_seen flag and improve test determinism Only drain remaining body after [DONE] is seen to avoid unnecessary work on early exit. Update test to explicitly close generator for deterministic finally block execution. --- tests/test_streaming.py | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/tests/test_streaming.py b/tests/test_streaming.py index 45691b58ce..6f5de94f0b 100644 --- a/tests/test_streaming.py +++ b/tests/test_streaming.py @@ -1,7 +1,7 @@ from __future__ import annotations -from collections.abc import Iterator, AsyncIterator from typing import Any +from collections.abc import Iterator, AsyncIterator import httpx2 import pytest @@ -278,6 +278,7 @@ def body() -> Iterator[bytes]: # Patch drain_async_iterator to track if it's called import openai._streaming as streaming_module + original_drain = streaming_module.drain_async_iterator streaming_module.drain_async_iterator = patched_drain @@ -291,11 +292,13 @@ def body() -> Iterator[bytes]: # Trailing body should NOT be exhausted since _done_seen is False assert exhausted is False - # Delete stream to finalize the generator and trigger finally block - del stream + # Explicitly close the generator to trigger finally block + await stream._iterator.aclose() # type: ignore[attr-defined] # Drain should NOT have been called since we didn't see [DONE] assert drained is False + # Response should be closed + assert response.is_closed is True finally: # Restore original drain function streaming_module.drain_async_iterator = original_drain