fix(streaming): wrap mid-stream request errors

src/openai/_httpx2.py

@@ -99,6 +99,18 @@ def timeout_exceptions() -> tuple[type[httpx2.TimeoutException], ...]:
     return (httpx2.TimeoutException,) if module is None else (httpx2.TimeoutException, module.TimeoutException)
 
 
+def request_exceptions() -> tuple[type[Exception], ...]:
+    module = _loaded_legacy_httpx()
+    if module is None:
+        return (httpx2.RequestError,)
+
+    legacy_request_error = cast(
+        type[Exception],
+        getattr(module, "RequestError", httpx2.RequestError),
+    )
+    return (httpx2.RequestError, legacy_request_error)
+
+
 def status_exceptions() -> tuple[type[httpx2.HTTPStatusError], ...]:
     module = _loaded_legacy_httpx()
     return (httpx2.HTTPStatusError,) if module is None else (httpx2.HTTPStatusError, module.HTTPStatusError)

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 request_exceptions, timeout_exceptions
+from ._exceptions import APIError, APITimeoutError, APIConnectionError
 
 if TYPE_CHECKING:
     from ._client import OpenAI, AsyncOpenAI
@@ -60,7 +61,15 @@ def __stream__(self) -> Iterator[_T]:
         iterator = self._iter_events()
 
         try:
-            for sse in iterator:
+            while True:
+                try:
+                    sse = next(iterator)
+                except StopIteration:
+                    break
+                except timeout_exceptions() as err:
+                    raise APITimeoutError(request=response.request) from err
+                except request_exceptions() as err:
+                    raise APIConnectionError(request=response.request) from err
                 if sse.data.startswith("[DONE]"):
                     break
 
@@ -170,7 +179,15 @@ async def __stream__(self) -> AsyncIterator[_T]:
         iterator = self._iter_events()
 
         try:
-            async for sse in iterator:
+            while True:
+                try:
+                    sse = await iterator.__anext__()
+                except StopAsyncIteration:
+                    break
+                except timeout_exceptions() as err:
+                    raise APITimeoutError(request=response.request) from err
+                except request_exceptions() as err:
+                    raise APIConnectionError(request=response.request) from err
                 if sse.data.startswith("[DONE]"):
                     break
 

src/openai/lib/streaming/_assistants.py

@@ -7,10 +7,11 @@
 
 from ..._utils import is_dict, is_list, consume_sync_iterator, consume_async_iterator
 from ..._compat import model_dump
-from ..._httpx2 import timeout_exceptions
+from ..._httpx2 import request_exceptions, timeout_exceptions
 from ..._models import construct_type
 from ..._streaming import Stream, AsyncStream
 from ...types.beta import AssistantStreamEvent
+from ..._exceptions import APIConnectionError
 from ...types.beta.threads import (
     Run,
     Text,
@@ -28,6 +29,14 @@ def _timeout_exceptions() -> tuple[type[Exception], ...]:
     return (*timeout_exceptions(), asyncio.TimeoutError)
 
 
+def _request_error_from_api_error(error: APIConnectionError) -> Exception | None:
+    cause = error.__cause__
+    if isinstance(cause, request_exceptions()):
+        return cause
+
+    return None
+
+
 class AssistantEventHandler:
     text_deltas: Iterable[str]
     """Iterator over just the text deltas in the stream.
@@ -406,9 +415,20 @@ def __stream__(self) -> Iterator[AssistantStreamEvent]:
             raise RuntimeError("Stream has not been started yet")
 
         try:
-            for event in stream:
-                self._emit_sse_event(event)
+            while True:
+                try:
+                    event = next(stream)
+                except StopIteration:
+                    break
+                except APIConnectionError as exc:
+                    error = _request_error_from_api_error(exc)
+
+                    if error is None:
+                        raise
 
+                    raise error from None
+
+                self._emit_sse_event(event)
                 yield event
         except _timeout_exceptions() as exc:
             self.on_timeout()
@@ -838,9 +858,20 @@ async def __stream__(self) -> AsyncIterator[AssistantStreamEvent]:
             raise RuntimeError("Stream has not been started yet")
 
         try:
-            async for event in stream:
-                await self._emit_sse_event(event)
+            while True:
+                try:
+                    event = await stream.__anext__()
+                except StopAsyncIteration:
+                    break
+                except APIConnectionError as exc:
+                    error = _request_error_from_api_error(exc)
 
+                    if error is None:
+                        raise
+
+                    raise error from None
+
+                await self._emit_sse_event(event)
                 yield event
         except _timeout_exceptions() as exc:
             await self.on_timeout()

tests/test_httpx2.py

@@ -21,6 +21,7 @@
 from openai._response import StreamAlreadyConsumed
 from openai.providers import bedrock
 from openai._constants import DEFAULT_TIMEOUT
+from openai.types.beta import AssistantStreamEvent
 
 
 def model_list(request: httpx2.Request) -> httpx2.Response:
@@ -661,6 +662,129 @@ async def async_response(request: httpx2.Request) -> httpx2.Response:
     assert isinstance(async_handler.exception, httpx2.ReadTimeout)
 
 
+@pytest.mark.filterwarnings(
+    "ignore:The Assistants API is deprecated in favor of the Responses API:DeprecationWarning"
+)
+@pytest.mark.asyncio
+@pytest.mark.parametrize("sync", [True, False], ids=["sync", "async"])
+@pytest.mark.parametrize("timeout", [True, False], ids=["timeout", "connection"])
+async def test_assistant_stream_hook_api_errors_are_not_unwrapped(sync: bool, timeout: bool) -> None:
+    hook_request = httpx2.Request("GET", "https://example.test/hook")
+
+    if timeout:
+        cause: httpx2.RequestError = httpx2.ReadTimeout("hook timeout")
+        hook_error: APIConnectionError = APITimeoutError(hook_request)
+    else:
+        cause = httpx2.RemoteProtocolError("hook connection error")
+        hook_error = APIConnectionError(request=hook_request)
+
+    hook_error.__cause__ = cause
+
+    class SyncHandler(openai.AssistantEventHandler):
+        def __init__(self) -> None:
+            super().__init__()
+            self.timed_out = False
+            self.exception: Exception | None = None
+
+        @override
+        def on_event(self, event: AssistantStreamEvent) -> None:
+            raise hook_error
+
+        @override
+        def on_timeout(self) -> None:
+            self.timed_out = True
+
+        @override
+        def on_exception(self, exception: Exception) -> None:
+            self.exception = exception
+
+    class AsyncHandler(openai.AsyncAssistantEventHandler):
+        def __init__(self) -> None:
+            super().__init__()
+            self.timed_out = False
+            self.exception: Exception | None = None
+
+        @override
+        async def on_event(self, event: AssistantStreamEvent) -> None:
+            raise hook_error
+
+        @override
+        async def on_timeout(self) -> None:
+            self.timed_out = True
+
+        @override
+        async def on_exception(self, exception: Exception) -> None:
+            self.exception = exception
+
+    content = b'event: thread.created\ndata: {"id":"thread_test","created_at":0,"object":"thread"}\n\n'
+
+    def sync_response(request: httpx2.Request) -> httpx2.Response:
+        return httpx2.Response(
+            200,
+            headers={"content-type": "text/event-stream"},
+            content=content,
+            request=request,
+        )
+
+    async def async_response(request: httpx2.Request) -> httpx2.Response:
+        return httpx2.Response(
+            200,
+            headers={"content-type": "text/event-stream"},
+            content=content,
+            request=request,
+        )
+
+    if sync:
+        handler = SyncHandler()
+
+        with OpenAI(
+            api_key="test",
+            base_url="https://example.test/v1",
+            http_client=openai.DefaultHttpx2Client(
+                transport=httpx2.MockTransport(sync_response),
+                trust_env=False,
+            ),
+            max_retries=0,
+        ) as client:
+            with client.beta.threads.runs.stream(  # pyright: ignore[reportDeprecated]
+                assistant_id="asst_test",
+                thread_id="thread_test",
+                event_handler=handler,
+            ) as stream:
+                with pytest.raises(APIConnectionError) as exc_info:
+                    stream.until_done()
+
+        assert exc_info.value is hook_error
+        assert handler.exception is hook_error
+        assert not handler.timed_out
+        assert hook_error.__cause__ is cause
+
+    else:
+        handler = AsyncHandler()
+
+        async with AsyncOpenAI(
+            api_key="test",
+            base_url="https://example.test/v1",
+            http_client=openai.DefaultAsyncHttpx2Client(
+                transport=httpx2.MockTransport(async_response),
+                trust_env=False,
+            ),
+            max_retries=0,
+        ) as client:
+            async with client.beta.threads.runs.stream(  # pyright: ignore[reportDeprecated]
+                assistant_id="asst_test",
+                thread_id="thread_test",
+                event_handler=handler,
+            ) as stream:
+                with pytest.raises(APIConnectionError) as exc_info:
+                    await stream.until_done()
+
+        assert exc_info.value is hook_error
+        assert handler.exception is hook_error
+        assert not handler.timed_out
+        assert hook_error.__cause__ is cause
+
+
 async def test_sigv4_provider_preserves_httpx2_family_and_rejects_one_shot_bodies() -> None:
     pytest.importorskip("botocore")