From 5dc9a3fc10c0f66ea7b704e15e1e175a8108d115 Mon Sep 17 00:00:00 2001 From: zerafachris Date: Thu, 16 Jul 2026 10:49:44 +0200 Subject: [PATCH 1/5] fix: drain remaining stream bytes after [DONE] event to prevent connection force-close --- src/openai/_streaming.py | 12 ++++++++++-- 1 file changed, 10 insertions(+), 2 deletions(-) diff --git a/src/openai/_streaming.py b/src/openai/_streaming.py index 45c13cc11d..a10acb59c7 100644 --- a/src/openai/_streaming.py +++ b/src/openai/_streaming.py @@ -106,7 +106,11 @@ def __stream__(self) -> Iterator[_T]: response=response, ) finally: - # Ensure the response is closed even if the consumer doesn't read all data + # Ensure the underlying connection can be cleanly returned to the pool + # by consuming any remaining bytes (e.g. the HTTP/1.1 chunked terminator) + # before closing the response. + for _ in iterator: + pass response.close() def __enter__(self) -> Self: @@ -216,7 +220,11 @@ async def __stream__(self) -> AsyncIterator[_T]: response=response, ) finally: - # Ensure the response is closed even if the consumer doesn't read all data + # Ensure the underlying connection can be cleanly returned to the pool + # by consuming any remaining bytes (e.g. the HTTP/1.1 chunked terminator) + # before closing the response. + async for _ in iterator: + pass await response.aclose() async def __aenter__(self) -> Self: From c789ceb105cd7784dbe9d88c03408f6859682f2c Mon Sep 17 00:00:00 2001 From: zerafachris Date: Sat, 18 Jul 2026 17:31:12 +0200 Subject: [PATCH 2/5] fix: only drain stream after normal [DONE] termination The drain added in the previous commit ran unconditionally in the `finally` block, so it also executed when the caller stopped before the `[DONE]` event -- e.g. breaking out of iteration and letting the stream be collected (GeneratorExit), an APIError, or async cancellation. In those cases it consumed the rest of the network stream instead of closing it, which for a long-running or stalled completion blocks until the server finishes and defeats the close-on-incomplete-consumption behaviour. Track whether `[DONE]` was actually observed and only drain on that path; otherwise close the response as before. Applied to both Stream and AsyncStream. Co-Authored-By: Claude Opus 4.8 (1M context) --- src/openai/_streaming.py | 34 ++++++++++----- tests/test_streaming.py | 93 +++++++++++++++++++++++++++++++++++++++- 2 files changed, 116 insertions(+), 11 deletions(-) diff --git a/src/openai/_streaming.py b/src/openai/_streaming.py index a10acb59c7..f615ad3593 100644 --- a/src/openai/_streaming.py +++ b/src/openai/_streaming.py @@ -57,10 +57,12 @@ def __stream__(self) -> Iterator[_T]: response = self.response process_data = self._client._process_response_data iterator = self._iter_events() + terminated = False try: for sse in iterator: if sse.data.startswith("[DONE]"): + terminated = True break # we have to special case the Assistants `thread.` events since we won't have an "event" key in the data @@ -106,11 +108,16 @@ def __stream__(self) -> Iterator[_T]: response=response, ) finally: - # Ensure the underlying connection can be cleanly returned to the pool - # by consuming any remaining bytes (e.g. the HTTP/1.1 chunked terminator) - # before closing the response. - for _ in iterator: - pass + # 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. + # + # 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. + if terminated: + for _ in iterator: + pass response.close() def __enter__(self) -> Self: @@ -171,10 +178,12 @@ async def __stream__(self) -> AsyncIterator[_T]: response = self.response process_data = self._client._process_response_data iterator = self._iter_events() + terminated = False try: async for sse in iterator: if sse.data.startswith("[DONE]"): + terminated = True break # we have to special case the Assistants `thread.` events since we won't have an "event" key in the data @@ -220,11 +229,16 @@ async def __stream__(self) -> AsyncIterator[_T]: response=response, ) finally: - # Ensure the underlying connection can be cleanly returned to the pool - # by consuming any remaining bytes (e.g. the HTTP/1.1 chunked terminator) - # before closing the response. - async for _ in iterator: - pass + # 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. + # + # 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. + if terminated: + async for _ in iterator: + pass await response.aclose() async def __aenter__(self) -> Self: diff --git a/tests/test_streaming.py b/tests/test_streaming.py index 04f8e51abd..272cd50a99 100644 --- a/tests/test_streaming.py +++ b/tests/test_streaming.py @@ -1,6 +1,6 @@ from __future__ import annotations -from typing import Iterator, AsyncIterator +from typing import Iterator, Generator, AsyncIterator, AsyncGenerator, cast import httpx import pytest @@ -216,6 +216,97 @@ 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_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 From f24d09a95e70deb079f3e96f18305e1323258e9f Mon Sep 17 00:00:00 2001 From: zerafachris Date: Sun, 19 Jul 2026 10:11:35 +0200 Subject: [PATCH 3/5] fix: don't fail a completed stream when the trailing drain errors If the connection drops after a valid [DONE] but before the trailing bytes arrive, the drain raised out of the finally block, replacing a successful result with a transport error and skipping response.close(). Treat the drain as best-effort cleanup. --- src/openai/_streaming.py | 16 ++++++++++++++-- tests/test_streaming.py | 27 +++++++++++++++++++++++++++ 2 files changed, 41 insertions(+), 2 deletions(-) diff --git a/src/openai/_streaming.py b/src/openai/_streaming.py index f615ad3593..02d3a826cc 100644 --- a/src/openai/_streaming.py +++ b/src/openai/_streaming.py @@ -116,7 +116,13 @@ def __stream__(self) -> Iterator[_T]: # or cancellation -- draining would block until the server finishes sending, # so close the response instead. if terminated: - for _ in iterator: + # 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 response.close() @@ -237,7 +243,13 @@ async def __stream__(self) -> AsyncIterator[_T]: # or cancellation -- draining would block until the server finishes sending, # so close the response instead. if terminated: - async for _ in iterator: + # 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 await response.aclose() diff --git a/tests/test_streaming.py b/tests/test_streaming.py index 272cd50a99..aa3a14f1a3 100644 --- a/tests/test_streaming.py +++ b/tests/test_streaming.py @@ -241,6 +241,33 @@ def body() -> Iterator[bytes]: 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_does_not_drain_on_premature_termination( From 04fdd7d7208e258c750eb5133b52af7c233faa1d Mon Sep 17 00:00:00 2001 From: zerafachris Date: Mon, 20 Jul 2026 09:09:27 +0200 Subject: [PATCH 4/5] fix: release the connection when the trailing drain is cancelled The best-effort drain after [DONE] only guarded against httpx.HTTPError, so a CancelledError raised while awaiting the trailing bytes (a caller timeout shorter than the HTTPX read timeout, or a server stalling after [DONE]) propagated out of the finally block and skipped response.close()/aclose(), leaking the connection. Wrap the drain so the close always runs. Co-Authored-By: Claude Opus 4.8 (1M context) --- src/openai/_streaming.py | 48 +++++++++++++++++++++++++--------------- tests/test_streaming.py | 30 +++++++++++++++++++++++++ 2 files changed, 60 insertions(+), 18 deletions(-) diff --git a/src/openai/_streaming.py b/src/openai/_streaming.py index 02d3a826cc..05bc273be1 100644 --- a/src/openai/_streaming.py +++ b/src/openai/_streaming.py @@ -115,16 +115,21 @@ def __stream__(self) -> Iterator[_T]: # 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. - if 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: + try: + if 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 - except httpx.HTTPError: - pass - response.close() + 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 @@ -242,16 +247,23 @@ async def __stream__(self) -> AsyncIterator[_T]: # 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. - if 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: + try: + if 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 - except httpx.HTTPError: - pass - await response.aclose() + 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 aa3a14f1a3..d9cbec60e3 100644 --- a/tests/test_streaming.py +++ b/tests/test_streaming.py @@ -1,5 +1,6 @@ from __future__ import annotations +import asyncio from typing import Iterator, Generator, AsyncIterator, AsyncGenerator, cast import httpx @@ -268,6 +269,35 @@ def body() -> Iterator[bytes]: 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_does_not_drain_on_premature_termination( From c8e39cbc4abf4aaeda3c55f187569d03731a1bda Mon Sep 17 00:00:00 2001 From: zerafachris Date: Mon, 20 Jul 2026 19:27:07 +0200 Subject: [PATCH 5/5] fix: bound the post-[DONE] drain so a held-open connection can't stall it If a server or proxy sends [DONE] but leaves the response open (or keeps sending data), the drain added in a previous commit would read until EOF, letting a logically completed stream hang for minutes or indefinitely. Wrap the raw byte iterator so that once [DONE] has been observed, only _MAX_POST_DONE_DRAIN_BYTES (64 KiB) more bytes are consumed before the drain stops on its own, same as reaching EOF. This has to bound the raw byte stream rather than the parsed SSE-event stream, since comment-only lines never surface as an event and would otherwise let the decoder keep pulling bytes internally without ever handing control back to the drain loop. Addresses PR review comment from chatgpt-codex-connector on src/openai/_streaming.py (id 3612568144, "Bound the post-DONE drain"). Prepared with AI assistance (Claude Code, Anthropic), reviewed for correctness before submission. --- src/openai/_streaming.py | 64 +++++++++++++++++++++++++++++++++++----- tests/test_streaming.py | 45 +++++++++++++++++++++++++++- 2 files changed, 100 insertions(+), 9 deletions(-) diff --git a/src/openai/_streaming.py b/src/openai/_streaming.py index 05bc273be1..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,19 +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() - terminated = False try: for sse in iterator: if sse.data.startswith("[DONE]"): - terminated = True + self._terminated = True break # we have to special case the Assistants `thread.` events since we won't have an "event" key in the data @@ -111,12 +136,15 @@ def __stream__(self) -> Iterator[_T]: # 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 terminated: + 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. @@ -157,6 +185,7 @@ class AsyncStream(Generic[_T]): response: httpx.Response _options: Optional[FinalRequestOptions] = None _decoder: SSEDecoder | SSEBytesDecoder + _terminated: bool = False def __init__( self, @@ -181,20 +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() - terminated = False try: async for sse in iterator: if sse.data.startswith("[DONE]"): - terminated = True + self._terminated = True break # we have to special case the Assistants `thread.` events since we won't have an "event" key in the data @@ -243,12 +288,15 @@ async def __stream__(self) -> AsyncIterator[_T]: # 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 terminated: + 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. diff --git a/tests/test_streaming.py b/tests/test_streaming.py index d9cbec60e3..7d9c669284 100644 --- a/tests/test_streaming.py +++ b/tests/test_streaming.py @@ -7,7 +7,7 @@ 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 @@ -298,6 +298,49 @@ def body() -> Iterator[bytes]: 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(