fix: drain remaining stream bytes after [DONE] event to prevent connection force-close#3510
fix: drain remaining stream bytes after [DONE] event to prevent connection force-close#3510zerafachris wants to merge 6 commits into
Conversation
…ction force-close
![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: 5dc9a3fc10
ℹ️ 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".
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) <noreply@anthropic.com>
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: c789ceb105
ℹ️ 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".
| if terminated: | ||
| for _ in iterator: | ||
| pass | ||
| response.close() |
There was a problem hiding this comment.
If the server has already sent a valid [DONE] event but closes/reset the connection before its final transport bytes arrive, advancing iterator here can raise an httpx transport/protocol error. That turns a semantically completed stream into a failure during list(stream)/normal iteration, and because the exception exits this finally, response.close() is skipped as well. Drain failures after [DONE] should still close the response and not replace the successful completion (the async drain has the same issue).
Useful? React with 👍 / 👎.
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.
|
Good catch — this was a real bug in my change, fixed in f24d09a. The drain sat directly in the Once the server has sent if terminated:
try:
for _ in iterator:
pass
except httpx.HTTPError:
pass
response.close()Same shape on the async path. I scoped the Added Full Prepared with AI assistance (Claude Code, Anthropic), reviewed for correctness before submission. |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: f24d09a95e
ℹ️ 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".
| # 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: |
There was a problem hiding this comment.
When async iteration is cancelled after [DONE] but while this drain is awaiting the remaining response bytes (for example, a user wraps stream consumption in an asyncio/AnyIO timeout shorter than the HTTPX read timeout, or the server stalls after sending [DONE]), the cancellation is not an httpx.HTTPError and control skips the subsequent await response.aclose(), leaving the response open. The close needs to be protected by a nested finally around the best-effort drain so cancellation still releases the connection.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Confirmed and fixed in 04fdd7d.
You are right that this is not hypothetical: CancelledError inherits from BaseException, so it walks straight past the except httpx.HTTPError around the drain, exits the finally, and response.close() / await response.aclose() never runs — the connection is leaked rather than returned to the pool. The same hole applies to any non-HTTPError escaping the drain, not just cancellation.
Fix is to nest the best-effort drain inside its own try, with the close in a finally, on both the sync and async paths. The drain still swallows httpx.HTTPError (so a dropped connection after a valid [DONE] does not fail an already-complete stream, per your earlier P1), but cancellation now propagates to the caller and releases the connection.
Added test_drain_cancellation_after_done_still_closes_response, parametrised sync/async, which raises CancelledError from the body after [DONE] and asserts stream.response.is_closed. Verified it fails on the previous commit (assert False where False = <Response [200 OK]>.is_closed, both params) and passes now. Full tests/test_streaming.py (28) and tests/test_client.py (174 passed, 2 skipped) green; ruff check + format clean.
The P1 thread above is addressed by f24d09a and can be resolved.
Prepared with AI assistance (Claude Code, Anthropic), reviewed for correctness before submission.
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) <noreply@anthropic.com>
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 04fdd7d720
ℹ️ 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".
| for _ in iterator: | ||
| pass |
There was a problem hiding this comment.
When the server or a proxy sends the [DONE] SSE frame but leaves the HTTP response open, this loop keeps reading until EOF before returning from an otherwise completed list(stream). With the SDK's default read timeout of 600s in src/openai/_constants.py, a logically completed stream can now hang for minutes (or indefinitely if bytes keep arriving); the async drain below has the same behavior. Consider making this drain bounded/best-effort so completion after [DONE] is not blocked on EOF.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Good catch — fixed in c8e39cb.
The drain now stops after consuming up to 64 KiB of trailing bytes once [DONE] has been observed, instead of reading to EOF. Once that cap is hit, the drain generator ends on its own (same as reaching EOF), the outer for _ in iterator: pass loop exits normally, and response.close()/aclose() still runs.
One detail worth calling out: the cap had to be applied to the raw byte stream (response.iter_bytes()/aiter_bytes()), not the parsed SSE-event stream. Lines that never decode into an event (e.g. SSE comment lines starting with :) don't get surfaced to the for _ in iterator: pass loop at all — the decoder keeps pulling and discarding raw chunks internally without yielding control back. So a bound at the event level wouldn't have covered the "proxy keeps sending data" case you flagged; only a bound at the raw-byte level does. Added test_post_done_drain_is_bounded in tests/test_streaming.py, which reproduces exactly that scenario (comment-only trailing chunks) and fails against the prior code (drains the full payload) and passes with this change.
This still relies on the client's configured read timeout to bound a single stalled read when the server sends nothing further at all after [DONE] — fully interrupting an in-flight blocking socket read would need lower-level transport/cancellation changes, which felt out of scope for this fix.
Prepared with AI assistance (Claude Code, Anthropic), reviewed for correctness before submission.
…l 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.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: c8e39cbc4a
ℹ️ 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".
| # 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: |
There was a problem hiding this comment.
Fresh evidence in the current patch is that the byte cap is only checked inside _iter_raw_bytes() after response.iter_bytes() yields a chunk, but this drain still advances iterator once after [DONE]. If an upstream/proxy sends the [DONE] event and then leaves the HTTP response idle without sending more body bytes or EOF, list(stream) still blocks here until the HTTP read timeout/EOF, so the held-open connection hang remains for idle sockets; the async drain has the same issue at its async for loop.
Useful? React with 👍 / 👎.
1 participant
Problem
When streaming SSE responses, after receiving
data: [DONE], theStream.__stream__method immediately callsresponse.close()without draining remaining bytes from the iterator. This means the HTTP/1.1 chunked terminator (0\r\n\r\n) may still be buffered, causing httpcore to destroy the connection instead of gracefully returning it to the pool.This is a regression from commit
6132922cwhich removed the drain logic. Thefinallyblock callsresponse.close()without first consuming remaining bytes.Fix
In the
finallyblock of bothStream.__stream__andAsyncStream.__stream__, before callingresponse.close()/await response.aclose(), drain any remaining bytes from the iterator to ensure the chunked transfer completes cleanly.Changes
src/openai/_streaming.py: Drainiteratorin finally block for syncStream.__stream__(addfor _ in iterator: pass)src/openai/_streaming.py: Drainiteratorin finally block for asyncAsyncStream.__stream__(addasync for _ in iterator: pass)