diff --git a/src/openai/_streaming.py b/src/openai/_streaming.py index 17c4c55eb3..5660f39a06 100644 --- a/src/openai/_streaming.py +++ b/src/openai/_streaming.py @@ -13,7 +13,7 @@ from ._utils import is_mapping, extract_type_var_from_base from ._httpx2 import request_exceptions, timeout_exceptions from ._exceptions import APIError, APITimeoutError, APIConnectionError -from ._streaming_cleanup import drain_response_async +from ._streaming_cleanup import close_owned, run_cleanup, aclose_owned, drain_response_async if TYPE_CHECKING: from ._client import OpenAI, AsyncOpenAI @@ -49,22 +49,39 @@ def __next__(self) -> _T: return self._iterator.__next__() def __iter__(self) -> Iterator[_T]: - for item in self._iterator: - yield item + return self._iterator def _iter_events(self) -> Iterator[ServerSentEvent]: + chunks: Optional[Iterator[bytes]] = None + events: Optional[Iterator[ServerSentEvent]] = None + failure: Optional[BaseException] = None try: - yield from self._decoder.iter_bytes(self.response.iter_bytes()) - except timeout_exceptions() as err: - raise APITimeoutError(request=self.response.request) from err - except request_exceptions() as err: - raise APIConnectionError(request=self.response.request) from err + try: + chunks = self.response.iter_bytes() + events = self._decoder.iter_bytes(chunks) + # Not `yield from`: it would close `events` itself, and the + # cleanup below must stay the only close. + for sse in events: + yield sse + except timeout_exceptions() as err: + raise APITimeoutError(request=self.response.request) from err + except request_exceptions() as err: + raise APIConnectionError(request=self.response.request) from err + except GeneratorExit: + raise # Normal closure, not a failure that should outrank cleanup errors. + except BaseException as err: + failure = err + raise + finally: + # Close the iterators opened here; stopping early does not close them. + close_owned(events, chunks, primary=failure) 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() + failure: Optional[BaseException] = None try: for sse in iterator: @@ -113,9 +130,15 @@ def __stream__(self) -> Iterator[_T]: cast_to=cast_to, response=response, ) + except GeneratorExit: + raise # Normal closure, not a failure that should outrank cleanup errors. + except BaseException as err: + failure = err + raise finally: - # Ensure the response is closed even if the consumer doesn't read all data - response.close() + # Close the event iterator explicitly; stopping at [DONE] leaves it suspended. + # Ensure the response is closed even if the consumer doesn't read all data. + close_owned(iterator, response, primary=failure) def __enter__(self) -> Self: return self @@ -126,7 +149,8 @@ def __exit__( exc: BaseException | None, exc_tb: TracebackType | None, ) -> None: - self.close() + # A cleanup error must not replace the exception leaving the block. + close_owned(self, primary=exc) def close(self) -> None: """ @@ -134,7 +158,10 @@ def close(self) -> None: Automatically called if the response body is read to completion. """ - self.response.close() + # Concurrent close does not cancel or join an active consumer; that consumer + # closes the iterator when it unwinds. + running = getattr(self._iterator, "gi_running", False) + close_owned(None if running else self._iterator, self.response) class AsyncStream(Generic[_T]): @@ -162,24 +189,38 @@ def __init__( async def __anext__(self) -> _T: return await self._iterator.__anext__() - async def __aiter__(self) -> AsyncIterator[_T]: - async for item in self._iterator: - yield item + def __aiter__(self) -> AsyncIterator[_T]: + return self._iterator async def _iter_events(self) -> AsyncIterator[ServerSentEvent]: + chunks: Optional[AsyncIterator[bytes]] = None + events: Optional[AsyncIterator[ServerSentEvent]] = None + failure: Optional[BaseException] = None try: - async for sse in self._decoder.aiter_bytes(self.response.aiter_bytes()): - yield sse - except timeout_exceptions() as err: - raise APITimeoutError(request=self.response.request) from err - except request_exceptions() as err: - raise APIConnectionError(request=self.response.request) from err + try: + chunks = self.response.aiter_bytes() + events = self._decoder.aiter_bytes(chunks) + async for sse in events: + yield sse + except timeout_exceptions() as err: + raise APITimeoutError(request=self.response.request) from err + except request_exceptions() as err: + raise APIConnectionError(request=self.response.request) from err + except GeneratorExit: + raise # Normal closure, not a failure that should outrank cleanup errors. + except BaseException as err: + failure = err + raise + finally: + # Close the iterators opened here; an interrupted `async for` does not. + await aclose_owned(events, chunks, primary=failure) 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() + failure: Optional[BaseException] = None try: async for sse in iterator: @@ -233,9 +274,17 @@ async def __stream__(self) -> AsyncIterator[_T]: cast_to=cast_to, response=response, ) + except GeneratorExit: + raise # Normal closure, not a failure that should outrank cleanup errors. + except BaseException as err: + failure = err + raise finally: - # Ensure the response is closed even if the consumer doesn't read all data - await response.aclose() + # Close the event iterator explicitly. Stopping at [DONE] leaves it and the + # httpx2 byte iterators beneath it suspended; left to the garbage collector, + # asyncio closes those nested generators concurrently. + # Ensure the response is closed even if the consumer doesn't read all data. + await aclose_owned(iterator, response, primary=failure) async def __aenter__(self) -> Self: return self @@ -246,7 +295,8 @@ async def __aexit__( exc: BaseException | None, exc_tb: TracebackType | None, ) -> None: - await self.close() + # A cleanup error must not replace the exception leaving the block. + await run_cleanup(self.close, primary=exc) async def close(self) -> None: """ @@ -254,7 +304,10 @@ async def close(self) -> None: Automatically called if the response body is read to completion. """ - await self.response.aclose() + # Concurrent close does not cancel or join an active consumer; that consumer + # closes the iterator when it unwinds. + running = getattr(self._iterator, "ag_running", False) + await aclose_owned(None if running else self._iterator, self.response) async def aclose(self) -> None: """Close the response and release the connection. Alias for `close()`.""" diff --git a/src/openai/_streaming_cleanup.py b/src/openai/_streaming_cleanup.py index d05b57ff04..8ec298bc8c 100644 --- a/src/openai/_streaming_cleanup.py +++ b/src/openai/_streaming_cleanup.py @@ -1,6 +1,6 @@ from __future__ import annotations -from typing import AsyncIterator, cast +from typing import Any, Callable, Optional, Awaitable, AsyncIterator, cast import anyio import httpx2 @@ -27,3 +27,83 @@ async def drain_response_async(response: httpx2.Response) -> None: # cancellation. Cancellation still propagates after the response closes. with anyio.CancelScope(shield=True): await response.aclose() + + +def _cleanup_error(primary: Optional[BaseException], failure: Optional[BaseException]) -> Optional[BaseException]: + """Pick the cleanup failure to raise, if any. + + A failure already propagating wins, except that a control-flow exception raised + during cleanup (for example cancellation) is never dropped in favor of an ordinary + error. + """ + if failure is None: + return None + if primary is None: + return failure + if not isinstance(failure, Exception) and isinstance(primary, Exception): + return failure + return None + + +def _keep(failure: Optional[BaseException], err: BaseException) -> BaseException: + """Keep the first cleanup failure, but let a later control-flow exception replace an ordinary one.""" + if failure is None or (isinstance(failure, Exception) and not isinstance(err, Exception)): + return err + return failure + + +def close_owned(*resources: object, primary: Optional[BaseException] = None) -> None: + """Close resources this caller owns, in order. + + Every resource gets its close attempt; objects without ``close()`` are skipped. + The first cleanup failure is raised after all attempts unless ``primary``, the + failure already propagating, takes precedence (see ``_cleanup_error``). + """ + failure: Optional[BaseException] = None + for resource in resources: + try: + close = getattr(resource, "close", None) + if close is not None: + close() + except BaseException as err: # re-raised below, after the remaining closes + failure = _keep(failure, err) + error = _cleanup_error(primary, failure) + if error is not None: + raise error + + +async def aclose_owned(*resources: object, primary: Optional[BaseException] = None) -> None: + """Async counterpart of ``close_owned``. + + The closes are shielded from AnyIO cancellation so an already-cancelled scope cannot + leave a resource half closed; cancellation still propagates afterwards. The shield + has no deadline, so a closer that never returns blocks the caller. + """ + failure: Optional[BaseException] = None + with anyio.CancelScope(shield=True): + for resource in resources: + try: + aclose = getattr(resource, "aclose", None) + if aclose is not None: + await aclose() + except BaseException as err: # re-raised below, after the remaining closes + failure = _keep(failure, err) + error = _cleanup_error(primary, failure) + if error is not None: + raise error + + +async def run_cleanup(cleanup: Callable[[], Awaitable[Any]], *, primary: Optional[BaseException] = None) -> None: + """Run one async cleanup callable with the shield and error precedence of ``aclose_owned``. + + For owners whose async close method is not named ``aclose``. + """ + failure: Optional[BaseException] = None + with anyio.CancelScope(shield=True): + try: + await cleanup() + except BaseException as err: # re-raised below according to precedence + failure = err + error = _cleanup_error(primary, failure) + if error is not None: + raise error diff --git a/src/openai/lib/beta/agents/_stream.py b/src/openai/lib/beta/agents/_stream.py index a2cf143905..c849e36559 100644 --- a/src/openai/lib/beta/agents/_stream.py +++ b/src/openai/lib/beta/agents/_stream.py @@ -6,6 +6,7 @@ from ._result import OutputT, AgentTurnResult, AgentTurnResultError, AgentOutputParseError, AgentTurnResultCollection from ...._compat import cached_property from ...._streaming import Stream, AsyncStream +from ...._streaming_cleanup import close_owned, aclose_owned from ...streaming.agents._dispatch import ToolDispatcher, AsyncToolDispatcher from ....types.beta.agent_session_event import AgentSessionEvent @@ -26,8 +27,10 @@ def _collection(self) -> AgentTurnResultCollection[OutputT]: @override def __stream__(self) -> Iterator[AgentSessionEvent]: + events = super().__stream__() + failure: BaseException | None = None try: - for event in super().__stream__(): + for event in events: self._collection.accept(event) pending = self._dispatcher.prepare(event) if self._dispatcher is not None else None yield event @@ -35,9 +38,15 @@ def __stream__(self) -> Iterator[AgentSessionEvent]: self._dispatcher.dispatch(pending) except Exception as error: self._collection.record_error(error) + failure = error + raise + except GeneratorExit: + raise # Normal closure, not a failure that should outrank cleanup errors. + except BaseException as error: + failure = error raise finally: - self.close() + close_owned(events, self, primary=failure) def with_result_collection(self) -> Self: """Beta: retain final messages for a result; call before consuming events.""" @@ -53,21 +62,27 @@ def get_final_result(self) -> AgentTurnResult[OutputT]: the cached result without consuming more events. """ collector = self._collection.enable() + failure: BaseException | None = None try: - collector.check_outcome(self._dispatcher.handlers if self._dispatcher is not None else ()) - if not collector.is_done(): - for _ in self: - collector.check_outcome(self._dispatcher.handlers if self._dispatcher is not None else ()) - if collector.is_done(): - break - return self._collection.result() - except (AgentTurnResultError, AgentOutputParseError): + try: + collector.check_outcome(self._dispatcher.handlers if self._dispatcher is not None else ()) + if not collector.is_done(): + for _ in self: + collector.check_outcome(self._dispatcher.handlers if self._dispatcher is not None else ()) + if collector.is_done(): + break + return self._collection.result() + except (AgentTurnResultError, AgentOutputParseError): + raise + except Exception as error: + collector.cause = error + raise collector.error("observation_failed") from error + except BaseException as error: + failure = error raise - except Exception as error: - collector.cause = error - raise collector.error("observation_failed") from error finally: - self.close() + # A cleanup error must not replace the result error being raised. + close_owned(self, primary=failure) class AsyncAgentSessionEventStream(AsyncStream[AgentSessionEvent], Generic[OutputT]): @@ -81,8 +96,10 @@ def _collection(self) -> AgentTurnResultCollection[OutputT]: @override async def __stream__(self) -> AsyncIterator[AgentSessionEvent]: + events = super().__stream__() + failure: BaseException | None = None try: - async for event in super().__stream__(): + async for event in events: self._collection.accept(event) pending = self._dispatcher.prepare(event) if self._dispatcher is not None else None yield event @@ -90,9 +107,15 @@ async def __stream__(self) -> AsyncIterator[AgentSessionEvent]: await self._dispatcher.dispatch(pending) except Exception as error: self._collection.record_error(error) + failure = error + raise + except GeneratorExit: + raise # Normal closure, not a failure that should outrank cleanup errors. + except BaseException as error: + failure = error raise finally: - await self.close() + await aclose_owned(events, self, primary=failure) def with_result_collection(self) -> Self: """Beta: retain final messages for a result; call before consuming events.""" @@ -106,18 +129,24 @@ async def get_final_result(self) -> AgentTurnResult[OutputT]: with_result_collection() before consuming events. """ collector = self._collection.enable() + failure: BaseException | None = None try: - collector.check_outcome(self._dispatcher.handlers if self._dispatcher is not None else ()) - if not collector.is_done(): - async for _ in self: - collector.check_outcome(self._dispatcher.handlers if self._dispatcher is not None else ()) - if collector.is_done(): - break - return self._collection.result() - except (AgentTurnResultError, AgentOutputParseError): + try: + collector.check_outcome(self._dispatcher.handlers if self._dispatcher is not None else ()) + if not collector.is_done(): + async for _ in self: + collector.check_outcome(self._dispatcher.handlers if self._dispatcher is not None else ()) + if collector.is_done(): + break + return self._collection.result() + except (AgentTurnResultError, AgentOutputParseError): + raise + except Exception as error: + collector.cause = error + raise collector.error("observation_failed") from error + except BaseException as error: + failure = error raise - except Exception as error: - collector.cause = error - raise collector.error("observation_failed") from error finally: - await self.close() + # A cleanup error must not replace the result error being raised. + await aclose_owned(self, primary=failure) diff --git a/src/openai/lib/streaming/agents/_streams.py b/src/openai/lib/streaming/agents/_streams.py index b222574295..d7fb89b99c 100644 --- a/src/openai/lib/streaming/agents/_streams.py +++ b/src/openai/lib/streaming/agents/_streams.py @@ -12,6 +12,7 @@ from ...._types import Omit, Headers, NotGiven, omit, not_given from ._dispatch import ToolDispatcher, AsyncToolDispatcher, _request_options from ...._streaming import Stream, AsyncStream +from ...._streaming_cleanup import close_owned, run_cleanup, aclose_owned from ...beta.agents._result import ( OutputT, AgentTurnResult, @@ -131,6 +132,7 @@ def __init__( self._iterator: Iterator[AgentSessionEvent] | None = None self._entered = False self._closed = False + self._deferred_cleanup_error: BaseException | None = None def __enter__(self) -> Self: if self._entered or self._closed: @@ -146,8 +148,8 @@ def __enter__(self) -> Self: self._sessions.events.create( self._session_id, events=[self._input], idempotency_key=self._idempotency_key, **self._options ) - except BaseException: - self.close() + except BaseException as error: + close_owned(self, primary=error) raise self._iterator = self._iterate() return self @@ -155,13 +157,16 @@ def __enter__(self) -> Self: def __exit__( self, exc_type: type[BaseException] | None, exc: BaseException | None, exc_tb: TracebackType | None ) -> None: - self.close() + if exc is not None: + self._deferred_cleanup_error = None # The propagating exception wins. + close_owned(self, primary=exc) def __iter__(self) -> Self: return self def __next__(self) -> AgentSessionEvent: if self._closed: + self._raise_deferred_cleanup_error() raise StopIteration if self._iterator is None: raise RuntimeError("Use sessions.stream in a with block") @@ -186,30 +191,50 @@ def get_final_result(self) -> AgentTurnResult[OutputT]: To iterate first, call with_result_collection() before consuming events. """ collector = self._collection.enable() + failure: BaseException | None = None try: - collector.check_outcome(self._handlers) - if not collector.is_done(): - for _ in self: - collector.check_outcome(self._handlers) - if collector.is_done(): - break - return self._collection.result() - except (AgentTurnResultError, AgentOutputParseError): + try: + collector.check_outcome(self._handlers) + if not collector.is_done(): + for _ in self: + collector.check_outcome(self._handlers) + if collector.is_done(): + break + result = self._collection.result() + except (AgentTurnResultError, AgentOutputParseError): + raise + except Exception as error: + self._collection.record_error(error) + raise collector.error("observation_failed") from error + # The outcome is established; a terminal cleanup failure may surface now. + self._raise_deferred_cleanup_error() + return result + except BaseException as error: + failure = error raise - except Exception as error: - self._collection.record_error(error) - raise collector.error("observation_failed") from error finally: - self.close() + # A cleanup error must not replace the result error being raised. + if failure is not None: + self._deferred_cleanup_error = None + close_owned(self, primary=failure) def close(self) -> None: """Close the event connection without cancelling the backend turn.""" self._closed = True - if self._stream is not None: - self._stream.close() + # Concurrent close does not cancel or join an active consumer; that consumer + # closes the iterator when it unwinds. + running = getattr(self._iterator, "gi_running", False) + close_owned(None if running else self._iterator, self._stream) + self._raise_deferred_cleanup_error() + + def _raise_deferred_cleanup_error(self) -> None: + error, self._deferred_cleanup_error = self._deferred_cleanup_error, None + if error is not None: + raise error def _iterate(self) -> Iterator[AgentSessionEvent]: assert self._stream is not None + failure: BaseException | None = None try: for event in self._stream: if not self._state.accept(event): @@ -217,19 +242,31 @@ def _iterate(self) -> Iterator[AgentSessionEvent]: self._collection.accept(event) terminal = self._state.terminal(event) if terminal: - self.close() + # Close now, but report a cleanup failure only after this terminal + # event's outcome has been evaluated. + try: + self.close() + except Exception as error: + self._deferred_cleanup_error = error pending = self._dispatcher.prepare(event) yield event if terminal or self._closed: + self._raise_deferred_cleanup_error() return if pending is not None: self._dispatcher.dispatch(pending) raise RuntimeError("Session event stream ended before the turn reached idle or failed") except Exception as error: self._collection.record_error(error) + failure = error + raise + except GeneratorExit: + raise # Normal closure, not a failure that should outrank cleanup errors. + except BaseException as error: + failure = error raise finally: - self.close() + close_owned(self, primary=failure) class AsyncAgentSessionStream(Generic[OutputT]): @@ -269,6 +306,7 @@ def __init__( self._iterator: AsyncIterator[AgentSessionEvent] | None = None self._entered = False self._closed = False + self._deferred_cleanup_error: BaseException | None = None async def __aenter__(self) -> Self: if self._entered or self._closed: @@ -284,8 +322,8 @@ async def __aenter__(self) -> Self: await self._sessions.events.create( self._session_id, events=[self._input], idempotency_key=self._idempotency_key, **self._options ) - except BaseException: - await self.close() + except BaseException as error: + await run_cleanup(self.close, primary=error) raise self._iterator = self._iterate() return self @@ -293,13 +331,16 @@ async def __aenter__(self) -> Self: async def __aexit__( self, exc_type: type[BaseException] | None, exc: BaseException | None, exc_tb: TracebackType | None ) -> None: - await self.close() + if exc is not None: + self._deferred_cleanup_error = None # The propagating exception wins. + await run_cleanup(self.close, primary=exc) def __aiter__(self) -> Self: return self async def __anext__(self) -> AgentSessionEvent: if self._closed: + self._raise_deferred_cleanup_error() raise StopAsyncIteration if self._iterator is None: raise RuntimeError("Use sessions.stream in an async with block") @@ -324,30 +365,50 @@ async def get_final_result(self) -> AgentTurnResult[OutputT]: To iterate first, call with_result_collection() before consuming events. """ collector = self._collection.enable() + failure: BaseException | None = None try: - collector.check_outcome(self._handlers) - if not collector.is_done(): - async for _ in self: - collector.check_outcome(self._handlers) - if collector.is_done(): - break - return self._collection.result() - except (AgentTurnResultError, AgentOutputParseError): + try: + collector.check_outcome(self._handlers) + if not collector.is_done(): + async for _ in self: + collector.check_outcome(self._handlers) + if collector.is_done(): + break + result = self._collection.result() + except (AgentTurnResultError, AgentOutputParseError): + raise + except Exception as error: + self._collection.record_error(error) + raise collector.error("observation_failed") from error + # The outcome is established; a terminal cleanup failure may surface now. + self._raise_deferred_cleanup_error() + return result + except BaseException as error: + failure = error raise - except Exception as error: - self._collection.record_error(error) - raise collector.error("observation_failed") from error finally: - await self.close() + # A cleanup error must not replace the result error being raised. + if failure is not None: + self._deferred_cleanup_error = None + await run_cleanup(self.close, primary=failure) async def close(self) -> None: """Close the event connection without cancelling the backend turn.""" self._closed = True - if self._stream is not None: - await self._stream.close() + # Concurrent close does not cancel or join an active consumer; that consumer + # closes the iterator when it unwinds. + running = getattr(self._iterator, "ag_running", False) + await aclose_owned(None if running else self._iterator, self._stream) + self._raise_deferred_cleanup_error() + + def _raise_deferred_cleanup_error(self) -> None: + error, self._deferred_cleanup_error = self._deferred_cleanup_error, None + if error is not None: + raise error async def _iterate(self) -> AsyncIterator[AgentSessionEvent]: assert self._stream is not None + failure: BaseException | None = None try: async for event in self._stream: if not self._state.accept(event): @@ -355,16 +416,28 @@ async def _iterate(self) -> AsyncIterator[AgentSessionEvent]: self._collection.accept(event) terminal = self._state.terminal(event) if terminal: - await self.close() + # Close now, but report a cleanup failure only after this terminal + # event's outcome has been evaluated. + try: + await self.close() + except Exception as error: + self._deferred_cleanup_error = error pending = self._dispatcher.prepare(event) yield event if terminal or self._closed: + self._raise_deferred_cleanup_error() return if pending is not None: await self._dispatcher.dispatch(pending) raise RuntimeError("Session event stream ended before the turn reached idle or failed") except Exception as error: self._collection.record_error(error) + failure = error + raise + except GeneratorExit: + raise # Normal closure, not a failure that should outrank cleanup errors. + except BaseException as error: + failure = error raise finally: - await self.close() + await run_cleanup(self.close, primary=failure) diff --git a/src/openai/lib/streaming/chat/_completions.py b/src/openai/lib/streaming/chat/_completions.py index 3476314f7f..374616f8bb 100644 --- a/src/openai/lib/streaming/chat/_completions.py +++ b/src/openai/lib/streaming/chat/_completions.py @@ -38,6 +38,7 @@ from ...._streaming import Stream, AsyncStream from ....types.chat import ChatCompletionChunk, ParsedChatCompletion, ChatCompletionToolUnionParam from ...._exceptions import LengthFinishReasonError, ContentFilterFinishReasonError +from ...._streaming_cleanup import close_owned, run_cleanup, aclose_owned from ....types.chat.chat_completion import Moderation, ChoiceLogprobs from ....types.chat.chat_completion_chunk import Choice as ChoiceChunk from ....types.chat.completion_create_params import ResponseFormat as ResponseFormatParam @@ -68,8 +69,7 @@ def __next__(self) -> ChatCompletionStreamEvent[ResponseFormatT]: return self._iterator.__next__() def __iter__(self) -> Iterator[ChatCompletionStreamEvent[ResponseFormatT]]: - for item in self._iterator: - yield item + return self._iterator def __enter__(self) -> Self: return self @@ -80,7 +80,8 @@ def __exit__( exc: BaseException | None, exc_tb: TracebackType | None, ) -> None: - self.close() + # A cleanup error must not replace the exception leaving the block. + close_owned(self, primary=exc) def close(self) -> None: """ @@ -88,7 +89,10 @@ def close(self) -> None: Automatically called if the response body is read to completion. """ - self._response.close() + # Concurrent close does not cancel or join an active consumer; that consumer + # closes the iterator when it unwinds. + running = getattr(self._iterator, "gi_running", False) + close_owned(None if running else self._iterator, self._raw_stream, self._response) def get_final_completion(self) -> ParsedChatCompletion[ResponseFormatT]: """Waits until the stream has been read to completion and returns @@ -163,7 +167,8 @@ def __exit__( exc_tb: TracebackType | None, ) -> None: if self.__stream is not None: - self.__stream.close() + # A cleanup error must not replace the exception leaving the block. + close_owned(self.__stream, primary=exc) class AsyncChatCompletionStream(Generic[ResponseFormatT]): @@ -190,9 +195,8 @@ def __init__( async def __anext__(self) -> ChatCompletionStreamEvent[ResponseFormatT]: return await self._iterator.__anext__() - async def __aiter__(self) -> AsyncIterator[ChatCompletionStreamEvent[ResponseFormatT]]: - async for item in self._iterator: - yield item + def __aiter__(self) -> AsyncIterator[ChatCompletionStreamEvent[ResponseFormatT]]: + return self._iterator async def __aenter__(self) -> Self: return self @@ -203,7 +207,8 @@ async def __aexit__( exc: BaseException | None, exc_tb: TracebackType | None, ) -> None: - await self.close() + # A cleanup error must not replace the exception leaving the block. + await run_cleanup(self.close, primary=exc) async def close(self) -> None: """ @@ -211,7 +216,10 @@ async def close(self) -> None: Automatically called if the response body is read to completion. """ - await self._response.aclose() + # Concurrent close does not cancel or join an active consumer; that consumer + # closes the iterator when it unwinds. + running = getattr(self._iterator, "ag_running", False) + await aclose_owned(None if running else self._iterator, self._raw_stream, self._response) async def get_final_completion(self) -> ParsedChatCompletion[ResponseFormatT]: """Waits until the stream has been read to completion and returns @@ -286,7 +294,8 @@ async def __aexit__( exc_tb: TracebackType | None, ) -> None: if self.__stream is not None: - await self.__stream.close() + # A cleanup error must not replace the exception leaving the block. + await run_cleanup(self.__stream.close, primary=exc) class ChatCompletionStreamState(Generic[ResponseFormatT]): diff --git a/src/openai/lib/streaming/responses/_responses.py b/src/openai/lib/streaming/responses/_responses.py index 6d379f8e6c..b05c3cc7f8 100644 --- a/src/openai/lib/streaming/responses/_responses.py +++ b/src/openai/lib/streaming/responses/_responses.py @@ -24,6 +24,7 @@ ResponseOutputItemDoneEvent, ResponseOutputItemAddedEvent, ) +from ...._streaming_cleanup import close_owned, run_cleanup, aclose_owned from ..._parsing._responses import TextFormatT, parse_text, parse_response from ....types.responses.tool_param import ToolParam from ....types.responses.parsed_response import ( @@ -54,8 +55,7 @@ def __next__(self) -> ResponseStreamEvent[TextFormatT]: return self._iterator.__next__() def __iter__(self) -> Iterator[ResponseStreamEvent[TextFormatT]]: - for item in self._iterator: - yield item + return self._iterator def __enter__(self) -> Self: return self @@ -73,7 +73,8 @@ def __exit__( exc: BaseException | None, exc_tb: TracebackType | None, ) -> None: - self.close() + # A cleanup error must not replace the exception leaving the block. + close_owned(self, primary=exc) def close(self) -> None: """ @@ -81,7 +82,10 @@ def close(self) -> None: Automatically called if the response body is read to completion. """ - self._response.close() + # Concurrent close does not cancel or join an active consumer; that consumer + # closes the iterator when it unwinds. + running = getattr(self._iterator, "gi_running", False) + close_owned(None if running else self._iterator, self._raw_stream, self._response) def get_final_response(self) -> ParsedResponse[TextFormatT]: """Waits until the stream has been read to completion and returns @@ -134,7 +138,8 @@ def __exit__( exc_tb: TracebackType | None, ) -> None: if self.__stream is not None: - self.__stream.close() + # A cleanup error must not replace the exception leaving the block. + close_owned(self.__stream, primary=exc) class AsyncResponseStream(Generic[TextFormatT]): @@ -155,9 +160,8 @@ def __init__( async def __anext__(self) -> ResponseStreamEvent[TextFormatT]: return await self._iterator.__anext__() - async def __aiter__(self) -> AsyncIterator[ResponseStreamEvent[TextFormatT]]: - async for item in self._iterator: - yield item + def __aiter__(self) -> AsyncIterator[ResponseStreamEvent[TextFormatT]]: + return self._iterator async def __stream__(self) -> AsyncIterator[ResponseStreamEvent[TextFormatT]]: async for sse_event in self._raw_stream: @@ -175,7 +179,8 @@ async def __aexit__( exc: BaseException | None, exc_tb: TracebackType | None, ) -> None: - await self.close() + # A cleanup error must not replace the exception leaving the block. + await run_cleanup(self.close, primary=exc) async def close(self) -> None: """ @@ -183,7 +188,10 @@ async def close(self) -> None: Automatically called if the response body is read to completion. """ - await self._response.aclose() + # Concurrent close does not cancel or join an active consumer; that consumer + # closes the iterator when it unwinds. + running = getattr(self._iterator, "ag_running", False) + await aclose_owned(None if running else self._iterator, self._raw_stream, self._response) async def get_final_response(self) -> ParsedResponse[TextFormatT]: """Waits until the stream has been read to completion and returns @@ -236,7 +244,8 @@ async def __aexit__( exc_tb: TracebackType | None, ) -> None: if self.__stream is not None: - await self.__stream.close() + # A cleanup error must not replace the exception leaving the block. + await run_cleanup(self.__stream.close, primary=exc) def _output_identity( diff --git a/tests/lib/streaming/agents/test_results.py b/tests/lib/streaming/agents/test_results.py index abca51d3b3..21bb779871 100644 --- a/tests/lib/streaming/agents/test_results.py +++ b/tests/lib/streaming/agents/test_results.py @@ -6,9 +6,11 @@ import httpx2 import pytest +from pydantic import BaseModel from openai import OpenAI, AsyncOpenAI -from openai.lib.beta.agents import AgentTurnResult, AgentTurnResultError +from openai._streaming import SSEDecoder +from openai.lib.beta.agents import AgentTurnResult, AgentTurnResultError, AgentOutputParseError from openai.lib.streaming.agents import AgentSessionStream, AsyncAgentSessionStream from tests.lib.streaming.agents.test_streams import Server, EventBody, sdk as sdk, call, idle, session, turn_event @@ -195,6 +197,183 @@ async def test_unsuccessful_turn(sdk: OpenAI | AsyncOpenAI, server: ResultServer assert server.body.closed +class _CloseFailingDecoder(SSEDecoder): + """Decodes normally; closing its event iterator fails.""" + + class _Events: + def __init__(self, inner: Any) -> None: + self._inner = inner + + def __iter__(self) -> Any: + return self + + def __next__(self) -> Any: + return next(self._inner) + + def __aiter__(self) -> Any: + return self + + async def __anext__(self) -> Any: + return await self._inner.__anext__() + + def close(self) -> None: + raise RuntimeError("decoder close failed") + + async def aclose(self) -> None: + raise RuntimeError("decoder close failed") + + @override + def iter_bytes(self, iterator: Iterator[bytes]) -> Any: + return self._Events(super().iter_bytes(iterator)) + + @override + def aiter_bytes(self, iterator: Any) -> Any: + return self._Events(super().aiter_bytes(iterator)) + + +async def test_cleanup_error_does_not_replace_result_error( + sdk: OpenAI | AsyncOpenAI, server: ResultServer, creation: bool, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(sdk, "_make_sse_decoder", lambda: _CloseFailingDecoder()) + server.body = EventBody([turn_event("created"), message(), turn_event("failed"), idle()]) + with pytest.raises(AgentTurnResultError) as exc: + await collect(sdk, creation) + assert exc.value.reason == "failed" + + +async def _final_result_through_stream(sdk: OpenAI | AsyncOpenAI, **kwargs: Any) -> Any: + if isinstance(sdk, AsyncOpenAI): + async with sdk.beta.agents.sessions.stream("session_test", input="Question", **kwargs) as async_stream: + return await async_stream.get_final_result() + with sdk.beta.agents.sessions.stream("session_test", input="Question", **kwargs) as stream: + return stream.get_final_result() + + +async def test_terminal_cleanup_error_does_not_replace_failed_session( + sdk: OpenAI | AsyncOpenAI, server: ResultServer, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(sdk, "_make_sse_decoder", lambda: _CloseFailingDecoder()) + server.body = EventBody( + [turn_event("created"), message(), {"type": "agent.session.failed", "session": session("failed")}] + ) + with pytest.raises(AgentTurnResultError) as exc: + await _final_result_through_stream(sdk) + assert exc.value.reason == "failed" + + +async def test_terminal_cleanup_error_does_not_replace_parse_error( + sdk: OpenAI | AsyncOpenAI, server: ResultServer, monkeypatch: pytest.MonkeyPatch +) -> None: + class Summary(BaseModel): + summary: str + + monkeypatch.setattr(sdk, "_make_sse_decoder", lambda: _CloseFailingDecoder()) + server.body = EventBody([turn_event("created"), message("not json"), turn_event("completed"), idle()]) + with pytest.raises(AgentOutputParseError): + await _final_result_through_stream(sdk, output_type=Summary) + + +async def test_terminal_cleanup_error_surfaces_after_successful_result( + sdk: OpenAI | AsyncOpenAI, server: ResultServer, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(sdk, "_make_sse_decoder", lambda: _CloseFailingDecoder()) + server.body = EventBody([turn_event("created"), message(), turn_event("completed"), idle()]) + with pytest.raises(RuntimeError, match="decoder close failed"): + await _final_result_through_stream(sdk) + + +@pytest.mark.parametrize("consume", ["iterate", "until_done", "close_after_terminal"]) +async def test_terminal_cleanup_error_surfaces_once_without_result( + sdk: OpenAI | AsyncOpenAI, server: ResultServer, monkeypatch: pytest.MonkeyPatch, consume: str +) -> None: + monkeypatch.setattr(sdk, "_make_sse_decoder", lambda: _CloseFailingDecoder()) + server.body = EventBody([turn_event("created"), message(), turn_event("completed"), idle()]) + raised = 0 + if isinstance(sdk, AsyncOpenAI): + async_stream = sdk.beta.agents.sessions.stream("session_test", input="Question") + try: + async with async_stream: + if consume == "iterate": + async for _ in async_stream: + pass + elif consume == "until_done": + await async_stream.until_done() + else: + async for event in async_stream: + if event.type == "agent.session.idle": + break + except RuntimeError as error: + assert str(error) == "decoder close failed" + raised += 1 + await async_stream.close() # Delivered once; a later close does not raise again. + else: + stream = sdk.beta.agents.sessions.stream("session_test", input="Question") + try: + with stream: + if consume == "iterate": + for _ in stream: + pass + elif consume == "until_done": + stream.until_done() + else: + for event in stream: + if event.type == "agent.session.idle": + break + except RuntimeError as error: + assert str(error) == "decoder close failed" + raised += 1 + stream.close() + assert raised == 1 + + +async def test_context_exception_wins_over_terminal_cleanup_error( + sdk: OpenAI | AsyncOpenAI, server: ResultServer, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(sdk, "_make_sse_decoder", lambda: _CloseFailingDecoder()) + server.body = EventBody([turn_event("created"), message(), turn_event("completed"), idle()]) + with pytest.raises(KeyError, match="caller"): + if isinstance(sdk, AsyncOpenAI): + async with sdk.beta.agents.sessions.stream("session_test", input="Question") as async_stream: + async for event in async_stream: + if event.type == "agent.session.idle": + raise KeyError("caller") + else: + with sdk.beta.agents.sessions.stream("session_test", input="Question") as stream: + for event in stream: + if event.type == "agent.session.idle": + raise KeyError("caller") + + +async def test_pre_terminal_context_exception_wins_over_cleanup_error( + sdk: OpenAI | AsyncOpenAI, server: ResultServer, creation: bool, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(sdk, "_make_sse_decoder", lambda: _CloseFailingDecoder()) + server.body = EventBody([turn_event("created"), message(), turn_event("completed"), idle()]) + with pytest.raises(KeyError, match="caller"): + if isinstance(sdk, AsyncOpenAI): + async_stream = ( + await sdk.beta.agents.sessions.create( + agent={"model": "test-model"}, environment={"type": "none"}, input="Question", stream=True + ) + if creation + else sdk.beta.agents.sessions.stream("session_test", input="Question") + ) + async with async_stream: + async for _ in async_stream: + raise KeyError("caller") + else: + stream = ( + sdk.beta.agents.sessions.create( + agent={"model": "test-model"}, environment={"type": "none"}, input="Question", stream=True + ) + if creation + else sdk.beta.agents.sessions.stream("session_test", input="Question") + ) + with stream: + for _ in stream: + raise KeyError("caller") + + @pytest.mark.parametrize("tail", [[], [turn_event("completed")]]) async def test_incomplete_answer( sdk: OpenAI | AsyncOpenAI, server: ResultServer, creation: bool, tail: list[dict[str, Any]] diff --git a/tests/test_streaming.py b/tests/test_streaming.py index 6378b1934d..06ae20d281 100644 --- a/tests/test_streaming.py +++ b/tests/test_streaming.py @@ -1,14 +1,16 @@ from __future__ import annotations +import gc import os import ssl +import sys import time import socket import asyncio import importlib import threading -from typing import Any, Iterator, AsyncIterator -from contextlib import aclosing, nullcontext +from typing import Any, Iterator, Generator, AsyncIterator +from contextlib import aclosing, nullcontext, contextmanager from http.server import ThreadingHTTPServer, BaseHTTPRequestHandler from typing_extensions import override @@ -18,6 +20,7 @@ from openai import OpenAI, AsyncOpenAI, APITimeoutError, APIConnectionError, DefaultAioHttpClient from openai._streaming import Stream, AsyncStream, ServerSentEvent +from openai.types.chat import ChatCompletionChunk @pytest.fixture( @@ -584,6 +587,310 @@ def body() -> Iterator[bytes]: assert response.is_closed +@contextmanager +def gc_finalized_async_generators() -> Generator[list[str], None, None]: + """Record async generators that reach the GC finalizer instead of being closed. + + Generators bind the finalizer on first iteration, so enter this before iterating. + """ + collected: list[str] = [] + firstiter, finalizer = sys.get_asyncgen_hooks() + + def record(agen: Any) -> None: + collected.append(agen.__qualname__) + if finalizer is not None: + finalizer(agen) + + sys.set_asyncgen_hooks(firstiter=firstiter, finalizer=record) + try: + yield collected + finally: + sys.set_asyncgen_hooks(firstiter=firstiter, finalizer=finalizer) + + +# Generators the SDK opens; httpx2-internal ones are owned by httpx2. +SDK_OWNED = ("AsyncStream.", "SSEDecoder.", "Response.aiter_bytes", "AsyncChatCompletionStream.") + + +def sse_body() -> Iterator[bytes]: + yield b'data: {"foo":true}\n\n' + yield b'data: {"bar":false}\n\n' + yield b"data: [DONE]\n\n" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("early_exit", [False, True], ids=["read-to-done", "early-exit"]) +async def test_async_stream_closes_its_iterators(early_exit: bool, async_client: AsyncOpenAI) -> None: + stream = AsyncStream( + cast_to=object, client=async_client, response=httpx2.Response(200, content=to_aiter(sse_body())) + ) + with gc_finalized_async_generators() as collected: + async for _ in stream: + if early_exit: + break + if early_exit: + await stream.close() + # Reading to [DONE] must clean up without an explicit close(). + del stream + gc.collect() + + assert [name for name in collected if name.startswith(SDK_OWNED)] == [] + + +@pytest.mark.asyncio +async def test_async_chat_completion_stream_closes_raw_stream_iterators(async_client: AsyncOpenAI) -> None: + from openai._types import omit + from openai.lib.streaming.chat import AsyncChatCompletionStream + + chunk = b'data: {"id":"c","object":"chat.completion.chunk","created":0,"model":"m","choices":[{"index":0,"delta":{"content":"hi"},"finish_reason":null}]}\n\n' + raw = AsyncStream( + cast_to=ChatCompletionChunk, + client=async_client, + response=httpx2.Response(200, content=to_aiter(iter([chunk, chunk, b"data: [DONE]\n\n"]))), + ) + stream = AsyncChatCompletionStream(raw_stream=raw, response_format=omit, input_tools=omit) + with gc_finalized_async_generators() as collected: + async for _ in stream: + break + await stream.close() + del stream, raw + gc.collect() + + assert [name for name in collected if name.startswith(SDK_OWNED)] == [] + + +class _TrackedEvents: + """A decoder iterator that is not a generator and must be closed exactly once.""" + + def __init__(self, fail_after: int | None = None, done_after: int | None = None, close_error: bool = False) -> None: + self.closes = 0 + self._sent = 0 + self._fail_after = fail_after + self._done_after = done_after + self._close_error = close_error + + def _next(self) -> ServerSentEvent: + if self._fail_after is not None and self._sent >= self._fail_after: + raise ValueError("decoder failed") + self._sent += 1 + if self._done_after is not None and self._sent > self._done_after: + return ServerSentEvent(data="[DONE]") + return ServerSentEvent(data='{"foo":true}') + + def __iter__(self) -> "_TrackedEvents": + return self + + def __next__(self) -> ServerSentEvent: + return self._next() + + def close(self) -> None: + self.closes += 1 + if self._close_error: + raise RuntimeError("close failed") + + def __aiter__(self) -> "_TrackedEvents": + return self + + async def __anext__(self) -> ServerSentEvent: + return self._next() + + async def aclose(self) -> None: + self.close() + + +class _TrackedDecoder: + def __init__(self, events: _TrackedEvents) -> None: + self.events = events + + def iter_bytes(self, iterator: Iterator[bytes]) -> Iterator[ServerSentEvent]: # noqa: ARG002 + return self.events + + def aiter_bytes(self, iterator: AsyncIterator[bytes]) -> AsyncIterator[ServerSentEvent]: # noqa: ARG002 + return self.events + + +@pytest.mark.asyncio +@pytest.mark.parametrize("sync", [True, False], ids=["sync", "async"]) +async def test_custom_decoder_iterator_is_closed_once(sync: bool, client: OpenAI, async_client: AsyncOpenAI) -> None: + events = _TrackedEvents() + if sync: + stream = Stream(cast_to=object, client=client, response=httpx2.Response(200, content=sse_body())) + stream._decoder = _TrackedDecoder(events) + next(iter(stream)) + stream.close() + else: + async_stream = AsyncStream( + cast_to=object, client=async_client, response=httpx2.Response(200, content=to_aiter(sse_body())) + ) + async_stream._decoder = _TrackedDecoder(events) + await async_stream.__anext__() + await async_stream.close() + + assert events.closes == 1 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("sync", [True, False], ids=["sync", "async"]) +async def test_cleanup_error_does_not_replace_stream_error( + sync: bool, client: OpenAI, async_client: AsyncOpenAI +) -> None: + events = _TrackedEvents(fail_after=1, close_error=True) + with pytest.raises(ValueError, match="decoder failed"): + if sync: + stream = Stream(cast_to=object, client=client, response=httpx2.Response(200, content=sse_body())) + stream._decoder = _TrackedDecoder(events) + for _ in stream: + pass + else: + async_stream = AsyncStream( + cast_to=object, client=async_client, response=httpx2.Response(200, content=to_aiter(sse_body())) + ) + async_stream._decoder = _TrackedDecoder(events) + async for _ in async_stream: + pass + + assert events.closes == 1 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("sync", [True, False], ids=["sync", "async"]) +@pytest.mark.parametrize("ending", ["done", "close"]) +async def test_cleanup_error_surfaces_without_other_error( + sync: bool, ending: str, client: OpenAI, async_client: AsyncOpenAI +) -> None: + # Normal closure (generator close() or [DONE]) must not hide a cleanup failure. + events = _TrackedEvents(done_after=1, close_error=True) + with pytest.raises(RuntimeError, match="close failed"): + if sync: + stream = Stream(cast_to=object, client=client, response=httpx2.Response(200, content=sse_body())) + stream._decoder = _TrackedDecoder(events) + if ending == "done": + for _ in stream: + pass + else: + next(iter(stream)) + stream.close() + else: + async_stream = AsyncStream( + cast_to=object, client=async_client, response=httpx2.Response(200, content=to_aiter(sse_body())) + ) + async_stream._decoder = _TrackedDecoder(events) + if ending == "done": + async for _ in async_stream: + pass + else: + await async_stream.__anext__() + await async_stream.close() + + assert events.closes == 1 + + +@pytest.mark.asyncio +async def test_aclose_owned_closes_every_resource_and_keeps_cancellation() -> None: + from openai._streaming_cleanup import aclose_owned + + closed: list[str] = [] + + class Closer: + def __init__(self, name: str, error: BaseException | None = None) -> None: + self.name = name + self.error = error + + async def aclose(self) -> None: + closed.append(self.name) + if self.error is not None: + raise self.error + + cancelled = asyncio.CancelledError() + with pytest.raises(asyncio.CancelledError): + await aclose_owned(Closer("first", cancelled), object(), Closer("second"), primary=ValueError("primary")) + + assert closed == ["first", "second"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("sync", [True, False], ids=["sync", "async"]) +@pytest.mark.parametrize("primary", [None, ValueError("primary")], ids=["no-primary", "ordinary-primary"]) +async def test_cleanup_keeps_later_cancellation_over_ordinary_error(sync: bool, primary: BaseException | None) -> None: + from openai._streaming_cleanup import close_owned, aclose_owned + + closed: list[str] = [] + + class Closer: + def __init__(self, name: str, error: BaseException | None = None) -> None: + self.name = name + self.error = error + + def close(self) -> None: + closed.append(self.name) + if self.error is not None: + raise self.error + + async def aclose(self) -> None: + self.close() + + resources = ( + Closer("ordinary", RuntimeError("close failed")), + Closer("cancelled", asyncio.CancelledError()), + Closer("last"), + ) + with pytest.raises(asyncio.CancelledError): + if sync: + close_owned(*resources, primary=primary) + else: + await aclose_owned(*resources, primary=primary) + + assert closed == ["ordinary", "cancelled", "last"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("sync", [True, False], ids=["sync", "async"]) +async def test_context_exit_keeps_caller_exception_over_cleanup_error( + sync: bool, client: OpenAI, async_client: AsyncOpenAI +) -> None: + events = _TrackedEvents(close_error=True) + with pytest.raises(KeyError, match="caller"): + if sync: + stream = Stream(cast_to=object, client=client, response=httpx2.Response(200, content=sse_body())) + stream._decoder = _TrackedDecoder(events) + with stream: + next(iter(stream)) + raise KeyError("caller") + else: + async_stream = AsyncStream( + cast_to=object, client=async_client, response=httpx2.Response(200, content=to_aiter(sse_body())) + ) + async_stream._decoder = _TrackedDecoder(events) + async with async_stream: + await async_stream.__anext__() + raise KeyError("caller") + assert events.closes == 1 + + +@pytest.mark.asyncio +async def test_context_exit_keeps_cancellation_over_cleanup_error(async_client: AsyncOpenAI) -> None: + events = _TrackedEvents(close_error=True) + started = asyncio.Event() + + async def consume() -> None: + async_stream = AsyncStream( + cast_to=object, client=async_client, response=httpx2.Response(200, content=to_aiter(sse_body())) + ) + async_stream._decoder = _TrackedDecoder(events) + async with async_stream: + await async_stream.__anext__() + started.set() + await asyncio.Event().wait() + + task = asyncio.create_task(consume()) + await started.wait() + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + assert task.cancelled() + assert events.closes == 1 + + async def to_aiter(iter: Iterator[bytes]) -> AsyncIterator[bytes]: for chunk in iter: yield chunk