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 5 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 👍 / 👎.
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)