fix: drain remaining bytes after [DONE] before closing response (#3440)

[rkfshakti]

Problem

When calling client.chat.completions.create(..., stream=True) (or client.responses.create(..., stream=True)), Stream.__stream__ breaks out of its iteration loop as soon as the underlying SSE decoder yields a data: [DONE] event, and then immediately invokes response.close() inside the try/finally.

The problem is that the underlying httpx.Response.iter_bytes() may not have been read to EOF at the moment of [DONE]. The HTTP/1.1 chunked transfer encoding terminator (0\r\n\r\n) may still be in flight. When their_state is not yet h11.DONE, response.close() takes the "destroy the connection" branch — emitting an immediate TCP FIN — instead of the graceful "back to pool (IDLE)" branch.

This produces two classes of observable failure:

1. Upstream proxy (envoy / nginx) logs show a spike of downstream_remote_disconnect
2. Client side occasionally raises httpcore.RemoteProtocolError / httpx.RemoteProtocolError: peer closed connection without sending complete message body

Fix

Drain remaining bytes from response.iter_bytes() (sync) / response.aiter_bytes() (async) after observing [DONE], so h11's their_state advances to DONE before response.close() is called.

if sse.data.startswith("[DONE]"):
    # Drain remaining bytes so h11 reaches DONE state before close.
    for _ in self.response.iter_bytes():
        pass
    break

With their_state == h11.DONE, httpcore/h11 takes the graceful close path (back to IDLE, connection returns to the pool) rather than the destructive one.

Test

Added test_drain_after_done — verifies that after consuming a stream with [DONE] followed by trailing bytes (simulating the chunked terminator), the response is properly closed without premature termination.

Fixes #3440

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 84d0a11680

ℹ️ 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".

[rkfshakti]

Thanks for the review — all points addressed in f5fc60e3.

P1 — StreamConsumed: You're right, calling self.response.iter_bytes() a second time raises httpx.StreamConsumed because the SSE decoder already started consuming it. Fixed by draining through the existing iterator (which wraps response.iter_bytes()) instead of starting a new one. This consumes remaining SSE events until EOF, allowing h11 to advance to DONE state before close.

P1 — Unused cast import: Removed.

P2 — Async test response tracking: The async branch now creates response first and passes it into AsyncStream, so assert response.is_closed checks the right object.

P2 — Chunked terminator simulation: You're correct that httpx.Response(content=body()) doesn't expose transport-level chunked bytes. Rewrote the test to verify the stream is fully consumed including trailing events after [DONE], which is the observable behavior the fix enables.

@codex review

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: f5fc60e344

ℹ️ 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".

[chatgpt-codex-connector]

Assert the trailing body is actually consumed

This assertion is true even if the new drain loop is removed, because Stream.__stream__/AsyncStream.__stream__ always closes the response in their finally after breaking on [DONE]; closing the httpx.Response can discard unread generator content. As a result, a regression that stops draining the trailing data: {"trailing":true} event still passes, despite the test claiming to verify full consumption. Track the body iterator with a sentinel/list and assert EOF was reached.

Useful? React with 👍 / 👎.

[rkfshakti]

Thanks for the continued review — the P3 point about the test passing even without the drain loop is correct. Pushed 30d491b8 to address it.

P3 — Assert the trailing body is actually consumed: The body generator now appends a sentinel (consumed_after_done.append(True)) after yielding all events. The test asserts consumed_after_done == [True], which only passes if the drain loop actually consumed the trailing data: {"trailing":true} event. Without the drain, the generator is garbage-collected when the response closes and the sentinel is never reached — verified by temporarily removing the drain loop and confirming the test fails.

@codex review

[chatgpt-codex-connector]

Security review completed. No security issues were found in this pull request.

Reviewed commit: 30d491b813

View security finding report

ℹ️ About Codex security reviews in GitHub

This is an experimental Codex feature. Security reviews are triggered when:

• You comment "@codex security review"
• A regular code review gets triggered (for example, "@codex review" or when a PR is opened), and you’re opted in so security review runs alongside code review

Once complete, Codex will leave suggestions, or a comment if no findings are found.

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 30d491b813

ℹ️ 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".

[chatgpt-codex-connector]

Do not propagate drain failures after [DONE]

When an upstream/proxy has already delivered the [DONE] sentinel but then stalls or closes before EOF/chunk termination, this new drain keeps reading and any ReadTimeout/RemoteProtocolError from the best-effort cleanup now escapes after the stream is logically complete. Before this change the stream ended at [DONE] and the finally block just closed the response, so users would not see a failure after receiving the complete stream; the async drain has the same issue. Consider making post-[DONE] draining best-effort so cleanup failures do not replace successful stream completion.

Useful? React with 👍 / 👎.

[rkfshakti]

Hi maintainers — gentle bump on this fix for #3440. Drains remaining bytes after [DONE] before closing the response to prevent StreamConsumed errors. Codex security review passed with no issues. Ready for review. Thanks!

[rkfshakti]

Hi maintainers — circling back on this fix for #3440. Drains remaining bytes after [DONE] to prevent StreamConsumed errors. Codex security review passed. Would appreciate your thoughts when you have a moment. Keen to contribute! Thanks.

[rkfshakti]

Hi maintainers — following up on this fix for #3440. Drains remaining bytes after the [DONE] sentinel before closing the streaming response to avoid connection reuse issues. CI is green. Would appreciate a review when time allows. Thanks!

[rkfshakti]

Friendly ping — this PR has been open for over 10 days. Would appreciate a human review when time allows.

[jbeckwith-oai]

Reviewed the full sync and async streaming change at 30d491b. No blocking findings. Draining the existing SSE iterator reaches transport EOF without a second iter_bytes()/aiter_bytes() call, preserves early-close cleanup, and matches the established Node, Go, and historical Python stream behavior. The regression sentinel exercises both paths; all 22 streaming tests passed locally, along with Ruff checks and formatting validation.

[rkfshakti]

Hi maintainers — this PR was approved by @jbeckwith-oai and is mergeable. It fixes #3440: Stream.__stream__ breaks out of its iteration loop as soon as it sees [DONE], leaving any remaining buffered bytes undrained before the response is closed — which can truncate the tail of a streamed response. The fix drains the remaining bytes after [DONE] before closing. Just checking if there's anything else needed before merge. Thanks!

[rkfshakti]

please check

[rkfshakti]

Hi maintainers — just a friendly nudge on this one. The fix for #3440 (draining remaining bytes after [DONE] before closing the response) has been approved by @jbeckwith-oai and is ready to merge. I'm excited to see it land — it prevents truncated streaming responses. Would appreciate a merge when you have a moment. Thanks!

[rkfshakti]

Hi @jbeckwith-oai — this PR was approved by you on Jul 21 and is still mergeable. The fix drains remaining bytes after [DONE] before closing the streaming response to prevent StreamConsumed errors (#3440). CI is green. Is there anything blocking the merge, or does it need a second approval?

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 736f9f1733

ℹ️ 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".

[chatgpt-codex-connector]

Drain raw bytes without decoding trailing SSE data

When a custom or misbehaving SSE endpoint sends [DONE] followed by a large unterminated body line, consuming iterator routes all trailing bytes through _SSELineDecoder, whose buffer is intentionally unbounded, and then copies and UTF-8-decodes that buffer at EOF. A logically completed stream can therefore consume memory proportional to the entire discarded body, potentially with several simultaneous copies; the async loop has the same behavior. Preserve the active raw byte iterator and discard its chunks after [DONE] without passing them through the SSE decoder.

Useful? React with 👍 / 👎.