fix(streaming): ensure remaining body is consumed after [DONE] in Str… - #3581
fix(streaming): ensure remaining body is consumed after [DONE] in Str…#3581vrs-darkness wants to merge 23 commits into
Conversation
…eam and AsyncStream - Added best-effort draining of remaining body bytes after receiving [DONE] to allow connection reuse. - Implemented error handling to prevent stream failures due to transport errors during draining. - Introduced tests to validate the behavior for both synchronous and asynchronous streams.
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: d4b47e7bb6
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
…and AsyncStream - Updated error handling in both Stream and AsyncStream classes to catch UnicodeError in addition to HTTPError during the draining of the stream after receiving [DONE]. - Added a new test to ensure that malformed trailing bytes after [DONE] do not cause failures in already-complete streams for both synchronous and asynchronous scenarios.
|
@codex review |
|
Codex Review: Didn't find any major issues. Swish! Reviewed commit: ℹ️ About Codex in GitHubCodex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
If Codex has suggestions, it will comment; otherwise it will react with 👍. When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback". |
|
Security review completed. No security issues were found in this pull request. Reviewed commit: Only the user who started this review can view the report in Codex. ℹ️ About Codex security reviews in GitHubThis is an experimental Codex feature. Security reviews are triggered when:
Once complete, Codex will leave suggestions, or a comment if no findings are found. |
|
@jbeckwith-oai please review PR whenever possible. It is currently a patch in my service would be great if you could check the PR and take it for a release which i could use directly. |
|
Hey Maintainers, friendly ping for reviewing the PR. Thanks! |
ting-hong-shieh
left a comment
ting-hong-shieh
left a comment
There was a problem hiding this comment.
Validated current head bc04f6d9 against its base 28888f9c and as a clean merge onto current main at d9029e3a (merge tree 7dc9e517). I still see one behavioral blocker and two readiness blockers.
The earlier bounded-drain concern still applies to this head. I ran the sync and async paths against the same in-process byte stream: one JSON event, [DONE], then a heartbeat held behind an explicit release event.
| Snapshot | Entered the trailing body | Returned before release |
|---|---|---|
Base 28888f9c |
no | yes |
Head bc04f6d9 |
yes | no |
Current-main merge tree 7dc9e517 |
yes | no |
Both changed snapshots return the same {"foo": true} chunk after the harness releases the body, and both close the response. Until that release, however, a completed stream waits for the trailing iterator. In a real transport this can delay completion until EOF or the read timeout; an unbounded custom transport can wait indefinitely.
The exact head's focused tests pass (26 passed), but its checked-in files do not currently pass the repository's pinned Ruff 0.14.7 gates: ruff check reports I001 for tests/test_streaming.py, and ruff format --check would reformat that file.
There is also a semantic main-branch integration failure that Git reports as a clean merge. main migrated these paths to httpx2, while the added drain blocks and tests still reference httpx. On merge tree 7dc9e517, the focused suite reports 20 passed, 6 failed with NameError: httpx is not defined; pinned Pyright 1.1.399 reports the same undefined name in both stream implementations and the new tests.
Please keep the post-[DONE] connection-reuse attempt bounded so [DONE] remains terminal for callers, and add a regression whose trailing iterator stays open. The branch also needs to be updated for the httpx2 migration and run through the repository's format/lint gates.
Evidence boundary: the behavioral snapshot uses only custom httpx byte streams and explicit synchronization events. It does not call an API, use credentials, or depend on an external server. The current-main result is from a synthetic clean merge tree, not a commit on this branch.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 73dae37905
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
…or bounded iterator handling - Updated the streaming logic to use `drain_sync_iterator` and `drain_async_iterator` instead of `consume_sync_iterator` and `consume_async_iterator`. - Enhanced comments to clarify the purpose of bounded draining after stream termination signals. - Added new drain functions to handle a limited number of items from iterators, preventing indefinite blocking. This change improves resource management and connection reuse after stream completion.
|
@codex review |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 114ec5741f
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
- Modified the streaming logic to utilize `iter_bytes` and `aiter_bytes` for both synchronous and asynchronous streams. - Enhanced comments to clarify the purpose of draining raw bytes after stream termination signals, ensuring proper resource management and preventing unbounded waits. This change improves the handling of byte streams in the streaming classes.
|
Security review completed. No security issues were found in this pull request. Reviewed commit: Only the user who started this review can view the report in Codex. ℹ️ About Codex security reviews in GitHubThis is an experimental Codex feature. Security reviews are triggered when:
Once complete, Codex will leave suggestions, or a comment if no findings are found. |
|
@codex review |
|
Security review completed. No security issues were found in this pull request. Reviewed commit: Only the user who started this review can view the report in Codex. ℹ️ About Codex security reviews in GitHubThis is an experimental Codex feature. Security reviews are triggered when:
Once complete, Codex will leave suggestions, or a comment if no findings are found. |
|
@ting-hong-shieh Behavioral blocker (bounded drain):
Regression test:
httpx2 migration:
CI verification:
|
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 8a8d02b20e
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
|
Thanks for the detailed update. I revalidated exact head The Two behavioral blockers still reproduce. First, the same sync/async held-open-body fixture emits one JSON event and
After release, both paths return the same Second, the custom-stream hook changes behavior on the same input. A sync and async
The direct decoder construction therefore bypasses the subclass override in both paths. The current custom-stream thread covers this regression. Both fresh threads are currently marked resolved on the unchanged head, but the conditions above remain reproducible. I have not added duplicate inline comments; my changes-requested review remains applicable until these two behaviors are addressed. Evidence boundary: both snapshots use in-process |
- Updated `drain_sync_iterator` and `drain_async_iterator` to implement timeout-based draining of remaining items from iterators, improving connection reuse after stream termination. - Enhanced comments to clarify the purpose of bounded draining and ensure proper resource management without indefinite blocking. This change optimizes the handling of byte streams in both synchronous and asynchronous contexts.
- Added a newline in the `drain_sync_iterator` function to improve code readability and maintain consistency in formatting. This change enhances the overall clarity of the code without altering functionality.
- Introduced `_byte_iterator` attributes in both `Stream` and `AsyncStream` classes to store byte iterators, enhancing the efficiency of event iteration. - Updated `_iter_events` methods to initialize `_byte_iterator` only when necessary, optimizing resource usage. - Adjusted `__stream__` methods to utilize the new `_iter_events` methods for better clarity and performance. This change refines the handling of byte streams, ensuring more efficient iteration and resource management.
- Removed unnecessary try-except blocks in `drain_sync_iterator` and `drain_async_iterator` functions to simplify error handling. - Enhanced the logic to ensure that exceptions are handled more cleanly, improving code readability and maintainability. This change refines the iterator draining process, ensuring clearer error management while maintaining functionality.
- Simplified comments in `Stream` and `AsyncStream` classes to clarify the purpose of draining remaining bytes for connection reuse. - Adjusted the timeout handling in `drain_async_iterator` to use the current running event loop, enhancing clarity and consistency. This change improves the readability of the code while maintaining the functionality of the iterator draining process.
- Updated `drain_sync_iterator` and `drain_async_iterator` functions to accept optional iterators, preventing indefinite blocking when a None value is passed. - Modified test cases to reflect changes in iterator behavior, ensuring that slow iterators trigger timeouts as expected. This change improves the robustness of iterator draining by handling None values gracefully and ensuring proper timeout behavior during draining.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: b94d8a3a2a
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
…ding - Modified `drain_sync_iterator` to run in a background thread, preventing blocking on `iterator.__next__()`. - Improved error handling within the draining process to ensure graceful termination on exceptions. This change optimizes the synchronous iterator draining process, enhancing responsiveness and robustness during iteration.
- Eliminated the unused `time` import from `_streams.py` to enhance code cleanliness and maintainability. This change contributes to a more streamlined codebase by removing unnecessary dependencies.
|
Hi @ting-hong-shieh, thanks for the detailed review! Both blockers you raised have been addressed in the latest commits:
Would appreciate if you could re-verify on the latest HEAD. Thanks! |
- Replaced asyncio-based timeout handling in `drain_async_iterator` with anyio's `move_on_after` for improved compatibility with both asyncio and Trio. - Streamlined the draining logic to utilize async iteration directly, enhancing code clarity and efficiency. This change enhances the flexibility of the async iterator draining process, ensuring better performance across different async backends.
ting-hong-shieh
left a comment
ting-hong-shieh
left a comment
There was a problem hiding this comment.
I revalidated exact head c1dc31b1 against base/main d9029e3a.
This revision fixes the earlier custom _iter_events() regression and both async timeout cases. The asyncio held-open-read and 10 ms heartbeat snapshots finish within 200 ms, and the same [DONE] fixture now succeeds under trio.run(), returning {"foo": true} and closing both the response and body.
One behavioral blocker remains in the synchronous background-thread implementation; I left the exact reproduction inline. One static gate also remains: Ruff 0.14.7 lint passes, but ruff format --check would reformat tests/test_streaming.py by adding a blank line after each local import in the two slow iterators.
Other exact-head checks pass: all 28 tests in tests/test_streaming.py, Pyright 1.1.399 with zero errors, MyPy 1.17 across 1,540 source files, and git diff --check.
Evidence boundary: the behavioral snapshots use only in-process httpx2 byte streams and explicit synchronization. They do not call an API, use credentials, or depend on an external server.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: c1dc31b190
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
- Updated `drain_sync_iterator` and `drain_async_iterator` to accept a `response` parameter, allowing for graceful closure of responses if draining times out. - Improved error handling during response closure to prevent unhandled exceptions, ensuring robust behavior in both synchronous and asynchronous contexts. This change enhances the reliability of iterator draining by ensuring that blocked reads are properly interrupted, improving overall stream management.
…] event - Introduced a `_done_seen` attribute in both `Stream` and `AsyncStream` classes to track the completion of the stream. - Updated the draining logic to only execute if the [DONE] event has been seen, preventing unnecessary consumption of data on early exits. - Enhanced comments to clarify the purpose of the changes in the draining process. This change optimizes resource management during stream termination, ensuring that abandoned data is not consumed unnecessarily.
…ning - Removed unnecessary line breaks in the function signatures of `drain_sync_iterator` and `drain_async_iterator` to improve code readability. - This change enhances the clarity of the code without altering the functionality of the iterator draining process.
…draining - Added a join operation for the background thread in `drain_sync_iterator` to ensure it exits gracefully after an interrupt. - This change enhances the robustness of the synchronous iterator draining process by providing adequate time for thread termination, improving overall resource management.
|
Hi @ting-hong-shieh, all your blockers are now addressed on the latest HEAD:
Would appreciate a re-review on the latest HEAD. Thanks again for the thorough feedback! |
ting-hong-shieh
left a comment
ting-hong-shieh
left a comment
There was a problem hiding this comment.
I revalidated exact head 5055298f after both worker-lifetime threads were resolved.
This revision closes the response after the first timeout and then joins the drain thread for another 50 ms. That bounds the caller’s additional wait, but it still does not guarantee worker termination. With the same held-open custom httpx2.SyncByteStream, the caller returns with response_closed_before_release=true and drain_still_blocked_after_return=true; the iterator exits only after the fixture releases it during cleanup. A custom next() that is not interrupted by response.close() can therefore still leave one daemon per completed stream.
Please keep this behavior open until a regression observes that the drain worker has terminated before fixture cleanup. Since a generic synchronous iterator has no cancellation primitive, the implementation may need to avoid starting an un-cancellable per-stream reader.
The rest of the revision is clean on this head. The async held-open and continuous-heartbeat cases, custom sync/async _iter_events() dispatch, and Trio snapshot pass. All 28 streaming tests, Ruff 0.14.7 lint and format, Pyright 1.1.399, MyPy 1.17 across 1,540 source files, and git diff --check also pass.
Evidence boundary: the snapshots use only in-process httpx2 byte streams and explicit synchronization. They do not call an API, use credentials, or depend on an external server.
|
Hi @ting-hong-shieh, you're right that the thread approach can't guarantee termination for custom transports. The The simple timeout loop alternative isn't great either — it can overshoot the deadline by the full duration of one blocking Given these constraints, the only honest fix for the sync path is to drop the drain entirely and revert to a plain Do you advise this approach — keeping the async drain as-is and simplifying sync back to plain response.close()? |
|
Yes — I’d take that trade-off. For the sync path, close the response directly after I would keep the async drain, with one wording caveat: Please split the current parameterized expectations accordingly:
That shape addresses the exact |
…ponse handling - Removed the `_done_seen` attribute from `Stream` and `AsyncStream` classes, streamlining the logic for handling the [DONE] event during iteration. - Updated the draining process to always close the response after the stream ends, ensuring proper resource management and connection reuse. - Adjusted comments for clarity regarding the purpose of draining remaining bytes. This change improves the efficiency and reliability of the streaming process by simplifying the code and enhancing response closure behavior.
|
Hi @ting-hong-shieh, updated on the latest HEAD 58d222d:
Ready for your snapshots. Thanks! |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 58d222d89b
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
- Removed the `_byte_iterator` attribute from the `Stream` class, streamlining the iteration logic. - Updated the `_iter_events` method to directly yield from the response's byte iterator, enhancing code clarity. - Added a comment to ensure the response is closed properly, improving resource management. This change simplifies the streaming process by reducing unnecessary state management and ensuring proper closure of resources.
ting-hong-shieh
left a comment
ting-hong-shieh
left a comment
There was a problem hiding this comment.
I revalidated exact head 990ed6fb against base/main d9029e3a.
The synchronous change fixes the worker-lifetime blocker. With the same held-open SyncByteStream, 5055298f entered the trailing body and returned with one live Thread-2 (_drain); 990ed6fb does not enter the trailing body, returns with no new live thread, and closes both the response and body.
One regression remains in the async path. This revision removes _done_seen from AsyncStream as well as Stream, so its finally block drains even when the caller exits before [DONE]. The same in-process early-exit fixture produces:
| Revision | Entered trailing body | Cancelled trailing read | Close time |
|---|---|---|---|
5055298f |
no | no | 0.0 ms |
990ed6fb |
yes | yes | 52.3 ms |
Both revisions return the same first event and close the response. The difference is that the new head resumes and consumes the abandoned body until move_on_after cancels it. Please retain _done_seen for AsyncStream and guard only the async drain with it; the sync path should remain a direct response.close().
The now-resolved bot suggestion to restore sync draining should not cause the background reader to return. Losing sync connection reuse is the deliberate trade-off that removes an un-cancellable per-stream worker.
All other exact-head checks pass: 28 streaming tests, Ruff 0.14.7 lint and format, Pyright 1.1.399 with zero errors, MyPy 1.17 across 1,540 source files, git diff --check, custom sync/async _iter_events() dispatch, held-open and heartbeat async deadlines, and the Trio cleanup snapshot.
Evidence boundary: the snapshots use only in-process httpx2 byte streams and explicit synchronization. They do not call an API, use credentials, or depend on an external server.
…ream - Introduced a `_done_seen` attribute in the `AsyncStream` class to track whether the [DONE] event has been encountered during streaming. - Updated the draining logic to only execute if `_done_seen` is true, preventing unnecessary consumption of trailing data when exiting early. - Added a new test to verify that early exits without seeing [DONE] do not trigger the drain function, ensuring efficient resource management. This change enhances the reliability of the streaming process by ensuring that only relevant data is consumed during termination.
2 participants
fix(streaming): best-effort drain after [DONE] for connection reuse
Changes being requested
When using a connection-pooled OpenAI client, each streaming request was adding ~300ms of overhead. SSE streams break on
data: [DONE]without fully reading the remaining response body; closing that partially-read response drops the connection from the pool, so the next request pays for a new TCP/TLS handshake.After
[DONE], drain the remaining iterator withconsume_sync_iterator/consume_async_iteratorbefore close so the body can be fully consumed and the pooled connection reused.Drain is best-effort:
httpx.HTTPErrorduring cleanup is swallowed so[DONE]stays terminal for callers even if a proxy leaves trailing heartbeats, the socket errors, or trailing bytes are truncated/malformed.