1313import time
1414from typing import List , Type , TypeVar , Callable , Optional , Awaitable
1515
16+ import httpx
17+
1618from .polling import PollingConfig , PollingTimeout
19+ from .._constants import LONG_POLL_CLIENT_BUFFER_SECONDS , STATUS_LONG_POLL_SERVER_MAX_SECONDS
1720from .._exceptions import APIStatusError , APIConnectionError
1821
1922T = TypeVar ("T" )
@@ -27,6 +30,7 @@ def wait_for_status(
2730 placeholder : Callable [[], T ],
2831 is_terminal : Callable [[T ], bool ],
2932 polling_config : Optional [PollingConfig ] = None ,
33+ server_max_timeout_seconds : float = STATUS_LONG_POLL_SERVER_MAX_SECONDS ,
3034) -> T :
3135 """Sync long-poll for a status change, retrying until *is_terminal* or timeout."""
3236 config = polling_config or PollingConfig ()
@@ -42,12 +46,18 @@ def wait_for_status(
4246 if remaining <= 0 :
4347 raise PollingTimeout (f"Exceeded timeout of { timeout } seconds" , last_result )
4448
49+ # Cap the server hold at what the server honors, and give the client read
50+ # timeout a buffer beyond it so the server returns 408 first.
51+ server_timeout = min (remaining , server_max_timeout_seconds )
4552 try :
4653 last_result = post_fn (
4754 path ,
48- body = {"statuses" : statuses , "timeout_seconds" : remaining },
55+ body = {"statuses" : statuses , "timeout_seconds" : server_timeout },
4956 cast_to = cast_to ,
50- options = {"max_retries" : 0 },
57+ options = {
58+ "max_retries" : 0 ,
59+ "timeout" : httpx .Timeout (server_timeout + LONG_POLL_CLIENT_BUFFER_SECONDS , connect = 5.0 ),
60+ },
5161 )
5262 except (APIConnectionError , APIStatusError ) as error :
5363 if isinstance (error , APIConnectionError ) or error .response .status_code == 408 :
@@ -67,6 +77,7 @@ async def async_wait_for_status(
6777 placeholder : Callable [[], T ],
6878 is_terminal : Callable [[T ], bool ],
6979 polling_config : Optional [PollingConfig ] = None ,
80+ server_max_timeout_seconds : float = STATUS_LONG_POLL_SERVER_MAX_SECONDS ,
7081) -> T :
7182 """Async long-poll for a status change, retrying until *is_terminal* or timeout."""
7283 config = polling_config or PollingConfig ()
@@ -82,12 +93,18 @@ async def async_wait_for_status(
8293 if remaining <= 0 :
8394 raise PollingTimeout (f"Exceeded timeout of { timeout } seconds" , last_result )
8495
96+ # Cap the server hold at what the server honors, and give the client read
97+ # timeout a buffer beyond it so the server returns 408 first.
98+ server_timeout = min (remaining , server_max_timeout_seconds )
8599 try :
86100 last_result = await post_fn (
87101 path ,
88- body = {"statuses" : statuses , "timeout_seconds" : remaining },
102+ body = {"statuses" : statuses , "timeout_seconds" : server_timeout },
89103 cast_to = cast_to ,
90- options = {"max_retries" : 0 },
104+ options = {
105+ "max_retries" : 0 ,
106+ "timeout" : httpx .Timeout (server_timeout + LONG_POLL_CLIENT_BUFFER_SECONDS , connect = 5.0 ),
107+ },
91108 )
92109 except (APIConnectionError , APIStatusError ) as error :
93110 if isinstance (error , APIConnectionError ) or error .response .status_code == 408 :
0 commit comments