22
33import os
44import ssl
5+ import time
6+ import socket
7+ import asyncio
58import importlib
9+ import threading
610from typing import Any , Iterator , AsyncIterator
711from contextlib import aclosing , nullcontext
12+ from http .server import ThreadingHTTPServer , BaseHTTPRequestHandler
13+ from typing_extensions import override
814
915import anyio
1016import httpx2
1117import pytest
1218
13- from openai import OpenAI , AsyncOpenAI , APITimeoutError , APIConnectionError
19+ from openai import OpenAI , AsyncOpenAI , APITimeoutError , APIConnectionError , DefaultAioHttpClient
1420from openai ._streaming import Stream , AsyncStream , ServerSentEvent
1521
1622
@@ -29,6 +35,167 @@ def http_module(request: pytest.FixtureRequest) -> Any:
2935 return importlib .import_module (request .param )
3036
3137
38+ @pytest .mark .parametrize ("mode" , ["async" , "aiohttp" ])
39+ @pytest .mark .parametrize ("ending" , ["complete" , "truncate" , "stall" , "trickle" ])
40+ async def test_done_http11_connection_reuse (mode : str , ending : str , http_module : Any ) -> None :
41+ if mode == "aiohttp" and http_module .__name__ == "httpx" :
42+ pytest .skip ("aiohttp uses HTTPX2" )
43+ connections = 0
44+ requests = 0
45+ release = threading .Event ()
46+
47+ class Handler (BaseHTTPRequestHandler ):
48+ protocol_version = "HTTP/1.1"
49+
50+ @override
51+ def setup (self ) -> None :
52+ nonlocal connections
53+ super ().setup ()
54+ connections += 1
55+ self .connection .setsockopt (socket .IPPROTO_TCP , socket .TCP_NODELAY , 1 )
56+
57+ def do_POST (self ) -> None :
58+ nonlocal requests
59+ self .rfile .read (int (self .headers ["Content-Length" ]))
60+ requests += 1
61+ self .send_response (200 )
62+ self .send_header ("Content-Type" , "text/event-stream" )
63+ self .send_header ("Transfer-Encoding" , "chunked" )
64+ self .end_headers ()
65+ payload = (
66+ b'data: {"id":"synthetic","object":"chat.completion.chunk","created":0,"model":"synthetic",'
67+ b'"choices":[{"index":0,"delta":{"content":"hello"}}]}\n \n data: [DONE]\n \n '
68+ )
69+ try :
70+ self .wfile .write (f"{ len (payload ):x} \r \n " .encode () + payload + b"\r \n " )
71+ if requests == 1 :
72+ if ending == "truncate" :
73+ self .close_connection = True
74+ return
75+ if ending == "stall" :
76+ release .wait (timeout = 0.6 )
77+ if ending == "trickle" :
78+ # Trailer bytes do not yield body chunks, but must still
79+ # obey the total cleanup deadline, not a per-read timeout.
80+ self .wfile .write (b"0\r \n X-Synthetic: " )
81+ for _ in range (15 ):
82+ if release .wait (timeout = 0.04 ):
83+ return
84+ self .wfile .write (b"a" )
85+ self .wfile .write (b"\r \n \r \n " )
86+ return
87+ self .wfile .write (b"0\r \n \r \n " )
88+ except ConnectionError :
89+ self .close_connection = True
90+
91+ @override
92+ def log_message (self , format : str , * args : object ) -> None : # noqa: A002
93+ pass
94+
95+ server = ThreadingHTTPServer (("127.0.0.1" , 0 ), Handler )
96+ thread = threading .Thread (target = server .serve_forever , kwargs = {"poll_interval" : 0.01 })
97+ thread .start ()
98+ options : dict [str , Any ] = {
99+ "api_key" : "synthetic" ,
100+ "base_url" : f"http://127.0.0.1:{ server .server_port } /v1" ,
101+ "max_retries" : 0 ,
102+ "timeout" : http_module .Timeout (5 , read = None ),
103+ }
104+ transport_options : dict [str , Any ] = {"trust_env" : False , "limits" : http_module .Limits (max_connections = 1 )}
105+ try :
106+ http_client = (
107+ DefaultAioHttpClient (** transport_options )
108+ if mode == "aiohttp"
109+ else http_module .AsyncClient (** transport_options )
110+ )
111+ async with AsyncOpenAI (** options , http_client = http_client ) as async_client :
112+ for _ in range (2 ):
113+ async with await async_client .chat .completions .create (
114+ model = "synthetic" , messages = [], stream = True
115+ ) as async_stream :
116+ assert (await async_stream .__anext__ ()).choices [0 ].delta .content == "hello"
117+ start = time .monotonic ()
118+ assert [chunk async for chunk in async_stream ] == []
119+ assert time .monotonic () - start < 0.4
120+ finally :
121+ release .set ()
122+ server .shutdown ()
123+ thread .join (timeout = 5 )
124+ server .server_close ()
125+
126+ assert not thread .is_alive ()
127+ assert requests == 2
128+ # A failed drain must discard the connection and release the only pool slot.
129+ assert connections == (1 if ending == "complete" else 2 )
130+
131+
132+ @pytest .mark .parametrize ("sync" , [True , False ], ids = ["sync" , "async" ])
133+ @pytest .mark .parametrize ("http_version" , [b"HTTP/1.1" , b"HTTP/2" ])
134+ async def test_done_discards_bytes_only_for_async_http11 (
135+ sync : bool , http_version : bytes , http_module : Any , client : OpenAI , async_client : AsyncOpenAI
136+ ) -> None :
137+ read_tail = False
138+
139+ def body () -> Iterator [bytes ]:
140+ nonlocal read_tail
141+ yield b"data: [DONE]\n \n "
142+ read_tail = True
143+ yield b"\xff \n \n " # Trailing data must not be decoded as SSE.
144+
145+ iterator = body ()
146+ async_iterator = to_aiter (iterator )
147+
148+ class Body (http_module .SyncByteStream , http_module .AsyncByteStream ):
149+ def __iter__ (self ) -> Iterator [bytes ]:
150+ return iterator
151+
152+ def __aiter__ (self ) -> AsyncIterator [bytes ]:
153+ return async_iterator
154+
155+ response = http_module .Response (200 , stream = Body (), extensions = {"http_version" : http_version })
156+ if sync :
157+ assert list (Stream (cast_to = object , response = response , client = client )) == []
158+ else :
159+ assert [event async for event in AsyncStream (cast_to = object , response = response , client = async_client )] == []
160+ assert response .is_closed
161+ assert read_tail == (not sync and http_version == b"HTTP/1.1" )
162+
163+
164+ @pytest .mark .parametrize ("cancel" , [True , False ], ids = ["caller-cancel" , "deadline" ])
165+ async def test_cancellation_during_done_drain_closes_response (
166+ cancel : bool , http_module : Any , async_client : AsyncOpenAI
167+ ) -> None :
168+ draining = asyncio .Event ()
169+
170+ class Body (http_module .AsyncByteStream ):
171+ closed = False
172+
173+ async def __aiter__ (self ) -> AsyncIterator [bytes ]:
174+ yield b"data: [DONE]\n \n "
175+ draining .set ()
176+ await asyncio .Event ().wait ()
177+
178+ async def aclose (self ) -> None :
179+ await anyio .sleep (0 )
180+ self .closed = True
181+
182+ body = Body ()
183+ response = http_module .Response (200 , stream = body )
184+ stream = AsyncStream (cast_to = object , response = response , client = async_client )
185+ task = asyncio .create_task (stream .__anext__ ())
186+ try :
187+ await asyncio .wait_for (draining .wait (), timeout = 5 )
188+ if cancel :
189+ task .cancel ()
190+ with pytest .raises (asyncio .CancelledError if cancel else StopAsyncIteration ):
191+ await task
192+ assert response .is_closed
193+ assert body .closed
194+ finally :
195+ task .cancel ()
196+ await asyncio .gather (task , return_exceptions = True )
197+
198+
32199@pytest .mark .parametrize ("sync" , [True , False ], ids = ["sync" , "async" ])
33200@pytest .mark .parametrize ("delivered" , [False , True ], ids = ["before-first-event" , "after-first-event" ])
34201@pytest .mark .parametrize (
0 commit comments