fix: drain remaining stream bytes after [DONE] event to prevent connection force-close

[zerafachris]

Problem

When streaming SSE responses, after receiving data: [DONE], the Stream.__stream__ method immediately calls response.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 6132922c which removed the drain logic. The finally block calls response.close() without first consuming remaining bytes.

Fix

In the finally block of both Stream.__stream__ and AsyncStream.__stream__, before calling response.close()/await response.aclose(), drain any remaining bytes from the iterator to ensure the chunked transfer completes cleanly.

Changes

• src/openai/_streaming.py: Drain iterator in finally block for sync Stream.__stream__ (add for _ in iterator: pass)
• src/openai/_streaming.py: Drain iterator in finally block for async AsyncStream.__stream__ (add async for _ in iterator: pass)

[chatgpt-codex-connector]

💡 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".

[chatgpt-codex-connector]

💡 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".

[chatgpt-codex-connector]

Preserve successful completion when trailing drain fails

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

[zerafachris]

Good catch — this was a real bug in my change, fixed in f24d09a.

The drain sat directly in the finally block, so a transport error while consuming the trailing bytes did two wrong things at once: it replaced an already-successful result with an exception, and because it propagated out of finally, it skipped response.close() / await response.aclose() entirely — leaking the connection in exactly the situation the original PR set out to fix.

Once the server has sent [DONE] the stream is semantically complete, so the drain is best-effort cleanup and should never be able to fail the call:

if terminated:
    try:
        for _ in iterator:
            pass
    except httpx.HTTPError:
        pass
response.close()

Same shape on the async path. I scoped the except to httpx.HTTPError rather than bare Exception so genuine programming errors still surface, and since CancelledError derives from BaseException on supported Pythons, cancellation still propagates correctly.

Added test_drain_failure_after_done_preserves_result (parametrized sync/async), which raises httpx.RemoteProtocolError from the response body after [DONE] and asserts both that the item is still returned and that response.is_closed is true. Verified it fails before the fix and passes after:

# before
FAILED tests/test_streaming.py::test_drain_failure_after_done_preserves_result[async]
FAILED tests/test_streaming.py::test_drain_failure_after_done_preserves_result[sync]
E   httpx.RemoteProtocolError: peer closed connection

# after
26 passed

Full tests/test_streaming.py passes (26), ruff check and ruff format --check clean.

Prepared with AI assistance (Claude Code, Anthropic), reviewed for correctness before submission.

[chatgpt-codex-connector]

💡 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".

[chatgpt-codex-connector]

Ensure async drain cancellation still closes response

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