diff --git a/src/openai/_streaming.py b/src/openai/_streaming.py index 45c13cc11d..d10b71afb4 100644 --- a/src/openai/_streaming.py +++ b/src/openai/_streaming.py @@ -19,6 +19,14 @@ _T = TypeVar("_T") +# After a normal `[DONE]` termination, only this many bytes of the response are +# consumed while draining. In the healthy case the trailing bytes are just +# transport framing (e.g. the HTTP/1.1 chunked terminator) -- a handful of bytes +# at most -- so this is generous headroom. It exists so that a server or proxy +# that leaves the connection open (or keeps sending data) after `[DONE]` can't +# turn "drain the last few bytes" into an unbounded read. +_MAX_POST_DONE_DRAIN_BYTES = 64 * 1024 + class Stream(Generic[_T]): """Provides the core interface to iterate over a synchronous stream response.""" @@ -26,6 +34,7 @@ class Stream(Generic[_T]): response: httpx.Response _options: Optional[FinalRequestOptions] = None _decoder: SSEBytesDecoder + _terminated: bool = False def __init__( self, @@ -50,17 +59,35 @@ def __iter__(self) -> Iterator[_T]: yield item def _iter_events(self) -> Iterator[ServerSentEvent]: - yield from self._decoder.iter_bytes(self.response.iter_bytes()) + yield from self._decoder.iter_bytes(self._iter_raw_bytes()) + + def _iter_raw_bytes(self) -> Iterator[bytes]: + """Wraps `response.iter_bytes()` so the post-`[DONE]` drain is bounded. + + Once `__stream__` sets `self._terminated` (having observed `[DONE]`), this + stops yielding new bytes after `_MAX_POST_DONE_DRAIN_BYTES` have been + consumed, so a connection that stays open -- or keeps sending data -- after + the stream has logically completed can't stall the drain indefinitely. + """ + drained = 0 + for chunk in self.response.iter_bytes(): + yield chunk + if self._terminated: + drained += len(chunk) + if drained >= _MAX_POST_DONE_DRAIN_BYTES: + return def __stream__(self) -> Iterator[_T]: cast_to = cast(Any, self._cast_to) response = self.response process_data = self._client._process_response_data + self._terminated = False iterator = self._iter_events() try: for sse in iterator: if sse.data.startswith("[DONE]"): + self._terminated = True break # we have to special case the Assistants `thread.` events since we won't have an "event" key in the data @@ -106,8 +133,31 @@ def __stream__(self) -> Iterator[_T]: response=response, ) finally: - # Ensure the response is closed even if the consumer doesn't read all data - response.close() + # Only drain when the stream terminated normally, i.e. we observed the + # `[DONE]` event and just need to consume the few remaining bytes (e.g. the + # HTTP/1.1 chunked terminator) so the connection can be returned to the pool. + # `_iter_raw_bytes()` bounds how much of that drain we're willing to do, so a + # connection that stays open (or keeps sending data) after `[DONE]` can't turn + # this into an unbounded read. + # + # On premature termination -- the caller breaking out of iteration, an error, + # or cancellation -- draining would block until the server finishes sending, + # so close the response instead. + try: + if self._terminated: + # Draining is best-effort cleanup for a stream that already completed. + # If the connection drops before the trailing bytes arrive, the result + # is still valid, so don't turn that into a failure for the caller. + try: + for _ in iterator: + pass + except httpx.HTTPError: + pass + finally: + # The drain can still be interrupted by something that isn't an + # `httpx.HTTPError` -- e.g. the caller cancelling iteration while we + # wait on the trailing bytes. The connection must be released either way. + response.close() def __enter__(self) -> Self: return self @@ -135,6 +185,7 @@ class AsyncStream(Generic[_T]): response: httpx.Response _options: Optional[FinalRequestOptions] = None _decoder: SSEDecoder | SSEBytesDecoder + _terminated: bool = False def __init__( self, @@ -159,18 +210,36 @@ 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()): + async for sse in self._decoder.aiter_bytes(self._aiter_raw_bytes()): yield sse + async def _aiter_raw_bytes(self) -> AsyncIterator[bytes]: + """Wraps `response.aiter_bytes()` so the post-`[DONE]` drain is bounded. + + Once `__stream__` sets `self._terminated` (having observed `[DONE]`), this + stops yielding new bytes after `_MAX_POST_DONE_DRAIN_BYTES` have been + consumed, so a connection that stays open -- or keeps sending data -- after + the stream has logically completed can't stall the drain indefinitely. + """ + drained = 0 + async for chunk in self.response.aiter_bytes(): + yield chunk + if self._terminated: + drained += len(chunk) + if drained >= _MAX_POST_DONE_DRAIN_BYTES: + return + async def __stream__(self) -> AsyncIterator[_T]: cast_to = cast(Any, self._cast_to) response = self.response process_data = self._client._process_response_data + self._terminated = False iterator = self._iter_events() try: async for sse in iterator: if sse.data.startswith("[DONE]"): + self._terminated = True break # we have to special case the Assistants `thread.` events since we won't have an "event" key in the data @@ -216,8 +285,33 @@ async def __stream__(self) -> AsyncIterator[_T]: response=response, ) finally: - # Ensure the response is closed even if the consumer doesn't read all data - await response.aclose() + # Only drain when the stream terminated normally, i.e. we observed the + # `[DONE]` event and just need to consume the few remaining bytes (e.g. the + # HTTP/1.1 chunked terminator) so the connection can be returned to the pool. + # `_aiter_raw_bytes()` bounds how much of that drain we're willing to do, so + # a connection that stays open (or keeps sending data) after `[DONE]` can't + # turn this into an unbounded read. + # + # On premature termination -- the caller breaking out of iteration, an error, + # or cancellation -- draining would block until the server finishes sending, + # so close the response instead. + try: + if self._terminated: + # Draining is best-effort cleanup for a stream that already completed. + # If the connection drops before the trailing bytes arrive, the result + # is still valid, so don't turn that into a failure for the caller. + try: + async for _ in iterator: + pass + except httpx.HTTPError: + pass + finally: + # The drain can still be interrupted by something that isn't an + # `httpx.HTTPError` -- e.g. `asyncio.CancelledError` when the caller + # wraps consumption in a timeout shorter than the HTTPX read timeout. + # `CancelledError` is a `BaseException`, so it bypasses the handler + # above; the connection must be released either way. + await response.aclose() async def __aenter__(self) -> Self: return self diff --git a/tests/test_streaming.py b/tests/test_streaming.py index 04f8e51abd..7d9c669284 100644 --- a/tests/test_streaming.py +++ b/tests/test_streaming.py @@ -1,12 +1,13 @@ from __future__ import annotations -from typing import Iterator, AsyncIterator +import asyncio +from typing import Iterator, Generator, AsyncIterator, AsyncGenerator, cast import httpx import pytest from openai import OpenAI, AsyncOpenAI -from openai._streaming import Stream, AsyncStream, ServerSentEvent +from openai._streaming import _MAX_POST_DONE_DRAIN_BYTES, Stream, AsyncStream, ServerSentEvent @pytest.mark.asyncio @@ -216,6 +217,196 @@ def body() -> Iterator[bytes]: assert sse.json() == {"content": "известни"} +@pytest.mark.asyncio +@pytest.mark.parametrize("sync", [True, False], ids=["sync", "async"]) +async def test_drains_remaining_bytes_after_done( + sync: bool, + client: OpenAI, + async_client: AsyncOpenAI, +) -> None: + """After a normal `[DONE]` termination the trailing bytes are consumed so that + the underlying connection can be returned to the pool.""" + consumed: list[bytes] = [] + + def body() -> Iterator[bytes]: + for chunk in [b'data: {"foo":true}\n\n', b"data: [DONE]\n\n", b"\n"]: + consumed.append(chunk) + yield chunk + + stream = make_stream(content=body(), sync=sync, client=client, async_client=async_client) + + items = await iter_all(stream) + assert len(items) == 1 + + # the whole body, including the bytes trailing `[DONE]`, was drained + assert consumed == [b'data: {"foo":true}\n\n', b"data: [DONE]\n\n", b"\n"] + + +@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: + """A transport error while draining must not fail an already-completed stream. + + The server sent a valid `[DONE]`, so the result is complete; a connection drop + while consuming the trailing bytes is cleanup noise and must still close the + response rather than propagating to the caller. + """ + + def body() -> Iterator[bytes]: + yield b'data: {"foo":true}\n\n' + yield b"data: [DONE]\n\n" + raise httpx.RemoteProtocolError("peer closed connection") + + stream = make_stream(content=body(), sync=sync, client=client, async_client=async_client) + + items = await iter_all(stream) + assert len(items) == 1 + + assert stream.response.is_closed + + +@pytest.mark.asyncio +@pytest.mark.parametrize("sync", [True, False], ids=["sync", "async"]) +async def test_drain_cancellation_after_done_still_closes_response( + sync: bool, + client: OpenAI, + async_client: AsyncOpenAI, +) -> None: + """Cancellation while draining must still release the connection. + + `[DONE]` has already been observed, so the drain is best-effort cleanup. If the + caller cancels while we wait on the trailing bytes -- e.g. an `asyncio` timeout + shorter than the HTTPX read timeout -- `CancelledError` is a `BaseException` and + so bypasses the `httpx.HTTPError` handler around the drain. The close has to sit + in a `finally` or the connection is leaked. + """ + + def body() -> Iterator[bytes]: + yield b'data: {"foo":true}\n\n' + yield b"data: [DONE]\n\n" + raise asyncio.CancelledError + + stream = make_stream(content=body(), sync=sync, client=client, async_client=async_client) + + with pytest.raises(asyncio.CancelledError): + await iter_all(stream) + + assert stream.response.is_closed + + +@pytest.mark.asyncio +@pytest.mark.parametrize("sync", [True, False], ids=["sync", "async"]) +async def test_post_done_drain_is_bounded( + sync: bool, + client: OpenAI, + async_client: AsyncOpenAI, +) -> None: + """The post-`[DONE]` drain must not turn into an unbounded read. + + A server or a proxy in between can send `[DONE]` but leave the HTTP response + open (or keep trickling bytes on it) instead of closing the connection. The + stream has already logically completed at that point, so draining should stop + after a bounded amount of data rather than consuming everything the other end + is willing to send. + """ + # Comment lines never decode into an `SSE` event, so none of these chunks make + # it back to the `for _ in iterator: pass` drain loop as a yielded event -- + # the only thing that can bound consumption here is the underlying byte read + # itself, not anything at the parsed-event level. + trailing_chunk = b": padding to simulate a proxy holding the connection open\n\n" + # comfortably more than the drain is allowed to consume + num_trailing_chunks = (_MAX_POST_DONE_DRAIN_BYTES // len(trailing_chunk)) * 4 + + consumed = 0 + + def body() -> Iterator[bytes]: + nonlocal consumed + yield b'data: {"foo":true}\n\n' + yield b"data: [DONE]\n\n" + for _ in range(num_trailing_chunks): + consumed += len(trailing_chunk) + yield trailing_chunk + + stream = make_stream(content=body(), sync=sync, client=client, async_client=async_client) + + items = await iter_all(stream) + assert len(items) == 1 + + # the drain stopped well short of consuming every trailing chunk the "server" sent + assert consumed < num_trailing_chunks * len(trailing_chunk) + assert stream.response.is_closed + + +@pytest.mark.asyncio +@pytest.mark.parametrize("sync", [True, False], ids=["sync", "async"]) +async def test_does_not_drain_on_premature_termination( + sync: bool, + client: OpenAI, + async_client: AsyncOpenAI, +) -> None: + """A caller that stops before `[DONE]` must not drain the rest of the stream. + + Draining here would block until the server finished sending, which for a + long-running or stalled completion defeats the point of breaking out early. + """ + trailing = 100 + consumed: list[bytes] = [] + + def body() -> Iterator[bytes]: + for chunk in [b'data: {"foo":true}\n\n', *([b'data: {"bar":true}\n\n'] * trailing), b"data: [DONE]\n\n"]: + consumed.append(chunk) + yield chunk + + stream = make_stream(content=body(), sync=sync, client=client, async_client=async_client) + + # consume a single event, then abandon iteration and let the stream be collected + await iter_next_item(stream) + await close_stream(stream) + + # the remaining events were left on the wire rather than being drained + assert len(consumed) < trailing, f"stream was drained: consumed {len(consumed)} chunks" + + +async def iter_all(stream: Stream[object] | AsyncStream[object]) -> list[object]: + if isinstance(stream, AsyncStream): + return [item async for item in stream] + + return list(stream) + + +async def iter_next_item(stream: Stream[object] | AsyncStream[object]) -> object: + if isinstance(stream, AsyncStream): + return await stream.__anext__() + + return next(iter(stream)) + + +async def close_stream(stream: Stream[object] | AsyncStream[object]) -> None: + """Close the underlying generator, mirroring what happens when an abandoned + stream object is garbage collected.""" + if isinstance(stream, AsyncStream): + await cast("AsyncGenerator[object, None]", stream._iterator).aclose() + else: + cast("Generator[object, None, None]", stream._iterator).close() + + +def make_stream( + content: Iterator[bytes], + *, + sync: bool, + client: OpenAI, + async_client: AsyncOpenAI, +) -> Stream[object] | AsyncStream[object]: + if sync: + return Stream(cast_to=object, client=client, response=httpx.Response(200, content=content)) + + return AsyncStream(cast_to=object, client=async_client, response=httpx.Response(200, content=to_aiter(content))) + + async def to_aiter(iter: Iterator[bytes]) -> AsyncIterator[bytes]: for chunk in iter: yield chunk