From 033cd40192ec2864eeedbbc9ee8068e676bd2bed Mon Sep 17 00:00:00 2001 From: Lazizbek Ergashev Date: Mon, 7 Sep 2026 19:36:06 +0500 Subject: [PATCH] fix: wrap transport failures while consuming a stream --- src/openai/_httpx2.py | 6 +++ src/openai/_streaming.py | 11 ++++- src/openai/lib/streaming/_assistants.py | 3 +- tests/test_httpx2.py | 12 +++-- tests/test_streaming.py | 58 ++++++++++++++++++++++++- 5 files changed, 83 insertions(+), 7 deletions(-) diff --git a/src/openai/_httpx2.py b/src/openai/_httpx2.py index 491398b43c..d65e10bd30 100644 --- a/src/openai/_httpx2.py +++ b/src/openai/_httpx2.py @@ -17,6 +17,7 @@ class _LegacyHttpxModule(Protocol): Timeout: type[httpx2.Timeout] Limits: type[httpx2.Limits] TimeoutException: type[httpx2.TimeoutException] + TransportError: type[httpx2.TransportError] HTTPStatusError: type[httpx2.HTTPStatusError] StreamConsumed: type[httpx2.StreamConsumed] RequestNotRead: type[httpx2.RequestNotRead] @@ -99,6 +100,11 @@ def timeout_exceptions() -> tuple[type[httpx2.TimeoutException], ...]: return (httpx2.TimeoutException,) if module is None else (httpx2.TimeoutException, module.TimeoutException) +def transport_exceptions() -> tuple[type[httpx2.TransportError], ...]: + module = _loaded_legacy_httpx() + return (httpx2.TransportError,) if module is None else (httpx2.TransportError, module.TransportError) + + def status_exceptions() -> tuple[type[httpx2.HTTPStatusError], ...]: module = _loaded_legacy_httpx() return (httpx2.HTTPStatusError,) if module is None else (httpx2.HTTPStatusError, module.HTTPStatusError) diff --git a/src/openai/_streaming.py b/src/openai/_streaming.py index 78e2d20aa7..f322c85fc9 100644 --- a/src/openai/_streaming.py +++ b/src/openai/_streaming.py @@ -11,7 +11,8 @@ import httpx2 from ._utils import is_mapping, extract_type_var_from_base -from ._exceptions import APIError +from ._httpx2 import timeout_exceptions, transport_exceptions +from ._exceptions import APIError, APITimeoutError, APIConnectionError if TYPE_CHECKING: from ._client import OpenAI, AsyncOpenAI @@ -106,6 +107,10 @@ def __stream__(self) -> Iterator[_T]: cast_to=cast_to, response=response, ) + except timeout_exceptions() as err: + raise APITimeoutError(request=response.request) from err + except transport_exceptions() as err: + raise APIConnectionError(request=response.request) from err finally: # Ensure the response is closed even if the consumer doesn't read all data response.close() @@ -216,6 +221,10 @@ async def __stream__(self) -> AsyncIterator[_T]: cast_to=cast_to, response=response, ) + except timeout_exceptions() as err: + raise APITimeoutError(request=response.request) from err + except transport_exceptions() as err: + raise APIConnectionError(request=response.request) from err finally: # Ensure the response is closed even if the consumer doesn't read all data await response.aclose() diff --git a/src/openai/lib/streaming/_assistants.py b/src/openai/lib/streaming/_assistants.py index 314961230d..9d5bd1fcc2 100644 --- a/src/openai/lib/streaming/_assistants.py +++ b/src/openai/lib/streaming/_assistants.py @@ -11,6 +11,7 @@ from ..._models import construct_type from ..._streaming import Stream, AsyncStream from ...types.beta import AssistantStreamEvent +from ..._exceptions import APITimeoutError from ...types.beta.threads import ( Run, Text, @@ -25,7 +26,7 @@ def _timeout_exceptions() -> tuple[type[Exception], ...]: - return (*timeout_exceptions(), asyncio.TimeoutError) + return (*timeout_exceptions(), asyncio.TimeoutError, APITimeoutError) class AssistantEventHandler: diff --git a/tests/test_httpx2.py b/tests/test_httpx2.py index 764fc00f0b..65f3a01823 100644 --- a/tests/test_httpx2.py +++ b/tests/test_httpx2.py @@ -638,11 +638,13 @@ async def async_response(request: httpx2.Request) -> httpx2.Response: with sync_client.beta.threads.runs.stream( # pyright: ignore[reportDeprecated] assistant_id="asst_test", thread_id="thread_test", event_handler=sync_handler ) as stream: - with pytest.raises(httpx2.ReadTimeout, match="assistant stream timeout"): + with pytest.raises(APITimeoutError) as sync_exc_info: stream.until_done() + assert isinstance(sync_exc_info.value.__cause__, httpx2.ReadTimeout) assert sync_handler.timed_out - assert isinstance(sync_handler.exception, httpx2.ReadTimeout) + assert isinstance(sync_handler.exception, APITimeoutError) + assert isinstance(sync_handler.exception.__cause__, httpx2.ReadTimeout) async_handler = AsyncHandler() async with AsyncOpenAI( @@ -654,11 +656,13 @@ async def async_response(request: httpx2.Request) -> httpx2.Response: async with async_client.beta.threads.runs.stream( # pyright: ignore[reportDeprecated] assistant_id="asst_test", thread_id="thread_test", event_handler=async_handler ) as async_stream: - with pytest.raises(httpx2.ReadTimeout, match="assistant stream timeout"): + with pytest.raises(APITimeoutError) as async_exc_info: await async_stream.until_done() + assert isinstance(async_exc_info.value.__cause__, httpx2.ReadTimeout) assert async_handler.timed_out - assert isinstance(async_handler.exception, httpx2.ReadTimeout) + assert isinstance(async_handler.exception, APITimeoutError) + assert isinstance(async_handler.exception.__cause__, httpx2.ReadTimeout) async def test_sigv4_provider_preserves_httpx2_family_and_rejects_one_shot_bodies() -> None: diff --git a/tests/test_streaming.py b/tests/test_streaming.py index ae6c0590f7..18729d40e9 100644 --- a/tests/test_streaming.py +++ b/tests/test_streaming.py @@ -5,7 +5,7 @@ import httpx2 import pytest -from openai import OpenAI, AsyncOpenAI +from openai import OpenAI, APIError, AsyncOpenAI, APITimeoutError, APIConnectionError from openai._streaming import Stream, AsyncStream, ServerSentEvent @@ -216,6 +216,35 @@ def body() -> Iterator[bytes]: assert sse.json() == {"content": "известни"} +@pytest.mark.parametrize("sync", [True, False], ids=["sync", "async"]) +@pytest.mark.parametrize( + "failure,expected", + [ + (httpx2.ReadTimeout("timed out while reading the stream"), APITimeoutError), + (httpx2.RemoteProtocolError("peer closed connection"), APIConnectionError), + ], + ids=["timeout", "connection"], +) +async def test_transport_error_mid_stream( + sync: bool, + failure: httpx2.TransportError, + expected: type[APIError], + client: OpenAI, + async_client: AsyncOpenAI, +) -> None: + def body() -> Iterator[bytes]: + yield b'data: {"foo":true}\n' + yield b"\n" + raise failure + + stream = make_stream(content=body(), sync=sync, client=client, async_client=async_client) + + with pytest.raises(expected) as exc_info: + await consume_stream(stream) + + assert exc_info.value.__cause__ is failure + + async def to_aiter(iter: Iterator[bytes]) -> AsyncIterator[bytes]: for chunk in iter: yield chunk @@ -246,3 +275,30 @@ def make_event_iterator( return AsyncStream( cast_to=object, client=async_client, response=httpx2.Response(200, content=to_aiter(content)) )._iter_events() + + +def make_stream( + content: Iterator[bytes], + *, + sync: bool, + client: OpenAI, + async_client: AsyncOpenAI, +) -> Stream[object] | AsyncStream[object]: + request = httpx2.Request("POST", "https://example.test/v1/chat/completions") + + if sync: + return Stream(cast_to=object, client=client, response=httpx2.Response(200, content=content, request=request)) + + return AsyncStream( + cast_to=object, client=async_client, response=httpx2.Response(200, content=to_aiter(content), request=request) + ) + + +async def consume_stream(stream: Stream[object] | AsyncStream[object]) -> None: + if isinstance(stream, AsyncStream): + async for _ in stream: + pass + return + + for _ in stream: + pass