Repository navigation
fix(streaming): close the SSE event iterators a stream opens - #1
Draft
fancyboi999 wants to merge 7 commits into
Draft
fancyboi999 wants to merge 7 commits into
fancyboi999 wants to merge 7 commits into
Conversation
AsyncStream stopped at [DONE] (or was closed early) without closing the async generators it had opened: _iter_events, the decoder iterator and the response.aiter_bytes() iterator, plus the __aiter__ wrapper after an early exit. They stayed suspended until the garbage collector reached them, where asyncio closes each one in its own aclose() task; a nested chain closed concurrently fails with "aclose(): asynchronous generator is already running" and skips its cleanup. Each layer now closes what it opened: _iter_events closes the decoder and byte iterators, __stream__ closes the event iterator before the response, close() closes the stream iterator when it is not running, and __aiter__ returns the stream iterator instead of wrapping it. The sync Stream gets the symmetric change.
- Keep the Iterator/AsyncIterator contracts so existing __stream__ overrides stay valid; close owned iterators through close_owned / aclose_owned, which skip objects without close()/aclose(). - Do not let a cleanup error replace the error or cancellation already propagating, and shield async cleanup from caller cancellation. - Create the byte and decoder iterators inside the exception-translation block, and loop explicitly in the sync stream so the decoder iterator is closed once. - close() closes the response even when iterator cleanup fails. - Chat and Responses stream wrappers close their own iterator and the raw stream instead of only the response; agent streams close the parent stream iterator they consume. - Tests cover the byte iterator, the chat wrapper, single close of custom decoder iterators, and error precedence.
- GeneratorExit from close() or [DONE] is normal closure; it no longer counts as a primary failure that silences cleanup errors. - close_owned/aclose_owned look up and call each closer inside the same handler, attempt every resource even after a BaseException, and never drop a cancellation raised during cleanup in favor of an ordinary error. - The response joins the same cleanup call as the iterators, so it is shielded and follows the same error precedence. - Comments no longer claim that closing the response ends a concurrent consumer. - Tests check [DONE] cleanup without an explicit close(), cleanup errors surfacing on normal closure, and aclose_owned continuing after cancellation.
…p errors - Cleanup helpers let a later control-flow exception (for example cancellation) replace an earlier ordinary cleanup failure. - Agent streams pass the propagating error into cleanup in get_final_result(), the event iterator and failed entry, so a decoder close failure no longer replaces AgentTurnResultError. - run_cleanup() covers owners whose async close method is not aclose(). - Tests cover cancellation after an ordinary cleanup error and a failed agent turn whose decoder close also fails.
AgentSessionStream closes its connection when it sees the terminal event, before yielding it. A cleanup failure there was raised before the caller evaluated the outcome, turning a failed turn or an output parse error into observation_failed. The failure is now kept and raised only after the outcome is established (or on the next read when iterating by hand); a result error still wins. close() also closes the stream's own iterator when it is not running.
The terminal close marks the stream closed, so the next read stopped before reaching the deferred cleanup failure and ordinary iteration, until_done() and close-after-terminal all hid it. Reads on a closed stream and close() now raise the pending failure once. A result error from get_final_result() or an exception leaving the with block still takes precedence and discards it.
…errors Context-manager exits closed streams without passing the exception that was leaving the block, so a cleanup failure could replace the caller's exception or a task cancellation. Stream, AsyncStream, the Chat and Responses streams and their managers, and the agent session streams now pass it as the primary failure.
1 task done
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Changes being requested
Streams now close the iterators they open instead of leaving them to the garbage collector. asyncio closes nested async generators that become unreachable together in concurrent
aclose()tasks. One of those tasks then fails withRuntimeError: aclose(): asynchronous generator is already running, logged asTask exception was never retrieved, and that generator's cleanup never runs. We see this in production with Python 3.13 and httpx2 2.13.1. The httpx2-internal part is pydantic/httpx2#1280, fixed by pydantic/httpx2#1281.Stream/AsyncStream_iterator,_iter_events, the decoder iterator and the byte iterator are each closed by the layer that opened it.close()also closes the stream iterator when it is not running.__iter__/__aiter__return the stream iterator directly instead of wrapping it in another generator.Iterator/AsyncIteratorcontracts are unchanged._streaming_cleanup.close_owned/aclose_owned/run_cleanup: one cleanup policy used everywhere.close/acloseare skipped.withblock) wins over a cleanup error.GeneratorExit,[DONE]) does not hide cleanup errors.Known limits, unchanged by this PR and documented in comments or docstrings:
close()does not cancel or join a consumer that is iterating concurrently.Additional context & links
Async generators finalized by the GC, from the repro in the issue (local SSE server, Python 3.13.16):
[DONE]close()_iter_events,SSEDecoder.aiter_bytes,Response.aiter_bytes,PoolByteStream.__aiter__AsyncStream.__aiter__PoolByteStream.__aiter__(httpx2-internal)PoolByteStream.__aiter__HTTP11ConnectionByteStream.__aiter__(httpcore2-internal)Tests, against
main718aee9 with Python 3.10 and the repo's uv lock:41 new test cases cover:
withexit.All fail on
mainor on the earlier commits of this branch where the gap existed, and pass at HEAD.Whole suite except
tests/api_resources(needs the Steady mock server) andtests/lib/live: 10201 passed on this branch and 10160 onmain. The same 200 failures and errors occur on both, from extras not installed here.ruff checkandruff formatare clean. mypy reports no issues. pyright 1.1.414 reports 0 errors.An independent model review (seven passes) ended with approve.
Not run:
tests/api_resources.