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