Skip to content

Commit a2753c7

Browse files
fix(streaming): leave request timeouts unchanged during bounded cleanup
1 parent 91c72d4 commit a2753c7

1 file changed

Lines changed: 5 additions & 14 deletions

File tree

‎src/openai/_streaming_cleanup.py‎

Lines changed: 5 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -22,10 +22,6 @@ def drain_response(response: httpx2.Response) -> None:
2222
if not isinstance(sock, socket.socket) or response.is_closed:
2323
return
2424

25-
saved_timeout = cast("dict[str, float | None] | None", response.request.extensions.get("timeout"))
26-
read_timeout = (saved_timeout or {}).get("read")
27-
timeout = _DRAIN_TIMEOUT if read_timeout is None else min(_DRAIN_TIMEOUT, read_timeout)
28-
2925
def interrupt_read() -> None:
3026
try:
3127
# Interrupt I/O without mutating SSLSocket's TLS state on this thread.
@@ -35,17 +31,16 @@ def interrupt_read() -> None:
3531

3632
# A read timeout alone is insufficient: HTTP trailers can trickle in without
3733
# yielding body chunks. Interrupt the socket when the total budget expires.
38-
deadline = time.monotonic() + timeout
39-
timer = threading.Timer(timeout, interrupt_read)
34+
deadline = time.monotonic() + _DRAIN_TIMEOUT
35+
timer = threading.Timer(_DRAIN_TIMEOUT, interrupt_read)
4036
timer.daemon = True
4137
try:
4238
timer.start()
4339
except RuntimeError:
4440
return # Cleanup is optional if a watchdog thread cannot be started.
4541
try:
46-
response.request.extensions["timeout"] = {**(saved_timeout or {}), "read": timeout}
47-
# Restart the transport iterator so httpcore picks up the cleanup timeout,
48-
# without passing discarded bytes through SSE or decompression again.
42+
# Drain below Response.iter_bytes() without decoding discarded bytes.
43+
# The watchdog supplies the deadline; the request's read timeout stays intact.
4944
for _ in cast(Iterable[bytes], response.stream):
5045
if time.monotonic() >= deadline:
5146
break
@@ -56,10 +51,6 @@ def interrupt_read() -> None:
5651
# The watchdog must finish before close() can return this connection to
5752
# the pool, or it could interrupt a subsequent request on the same socket.
5853
timer.join()
59-
if saved_timeout is None:
60-
response.request.extensions.pop("timeout", None)
61-
else:
62-
response.request.extensions["timeout"] = saved_timeout
6354

6455

6556
async def drain_response_async(response: httpx2.Response) -> None:
@@ -71,7 +62,7 @@ async def drain_response_async(response: httpx2.Response) -> None:
7162
# Also enforce the deadline on transports yielding buffered data.
7263
await checkpoint()
7364
except request_exceptions() + stream_consumed_exceptions():
74-
pass
65+
pass # Cleanup must not turn completed output into a request error.
7566
finally:
7667
# Close outside our timeout scope and shield cleanup from caller
7768
# cancellation. Cancellation still propagates after the response closes.

0 commit comments

Comments
 (0)