Coverage for src/forge_dashboard/retry.py: 100%

1"""Retry transport for httpx: exponential backoff with jitter on 429/5xx, 

2honoring a server-sent ``Retry-After`` (seconds or an HTTP-date) ahead of 

3the computed backoff. Never retries any other 4xx -- it won't succeed the 

4second time either, and retrying only delays the real error reaching the 

5caller (rules/sdk-generation.md's "Client shape"). 

6 

7httpx's transport-level hooks (``httpx.BaseTransport``/``AsyncBaseTransport``) 

8apply this uniformly to every request the client makes, rather than 

9wrapping each generated call site in a retry loop by hand. 

10""" 

11 

12from __future__ import annotations 

13 

14import asyncio 

15import random 

16import time 

17from dataclasses import dataclass 

18from datetime import UTC, datetime 

19from email.utils import parsedate_to_datetime 

20 

21import httpx 

22 

23_RETRYABLE_5XX_AND_429 = 500 

24 

25 

26@dataclass(frozen=True) 

27class RetryConfig: 

28 """Retry policy for :class:`RetryTransport`/:class:`AsyncRetryTransport`. 

29 

30 Attributes: 

31 max_retries: Number of additional attempts after the first. 

32 base_seconds: Starting backoff before jitter and any server-sent 

33 ``Retry-After``. 

34 """ 

35 

36 max_retries: int = 3 

37 base_seconds: float = 0.25 

38 

39 

40DEFAULT_RETRY_CONFIG = RetryConfig() 

41 

42 

43def _should_retry( 

44 attempt: int, 

45 config: RetryConfig, 

46 response: httpx.Response | None, 

47 exc: Exception | None, 

48) -> bool: 

49 if attempt >= config.max_retries: 

50 return False 

51 if exc is not None: 

52 # A network-level failure (connection reset, timeout) is 

53 # retry-worthy the same as a 5xx. 

54 return True 

55 if response is None: 

56 # Every call site passes exactly one of response/exc -- not an 

57 # assert, so this invariant still holds under -O and doesn't trip 

58 # bandit's B101. 

59 raise ValueError("_should_retry needs a response or an exception") 

60 return response.status_code == httpx.codes.TOO_MANY_REQUESTS or response.status_code >= _RETRYABLE_5XX_AND_429 

61 

62 

63def _retry_after_seconds(response: httpx.Response) -> float | None: 

64 value = response.headers.get("retry-after") # pragma: no mutate - httpx.Headers lookups are case-insensitive 

65 if not value: 

66 return None 

67 value = value.strip() 

68 try: 

69 return float(int(value)) 

70 except ValueError: 

71 pass 

72 try: 

73 when = parsedate_to_datetime(value) 

74 except (TypeError, ValueError): 

75 return None 

76 if when.tzinfo is None: 

77 when = when.replace(tzinfo=UTC) 

78 return max((when - datetime.now(UTC)).total_seconds(), 0.0) 

79 

80 

81def _backoff_seconds(attempt: int, config: RetryConfig) -> float: 

82 backoff = config.base_seconds * (2.0**attempt) 

83 # Retry jitter, not security-sensitive -- a predictable PRNG is fine. 

84 jitter = random.uniform(0, backoff) # noqa: S311 # nosec B311 

85 return backoff / 2 + jitter / 2 

86 

87 

88def _retry_delay(attempt: int, config: RetryConfig, response: httpx.Response | None) -> float: 

89 if response is not None: 

90 retry_after = _retry_after_seconds(response) 

91 if retry_after is not None: 

92 return retry_after 

93 return _backoff_seconds(attempt, config) 

94 

95 

96class RetryTransport(httpx.BaseTransport): 

97 """Wraps a sync ``httpx.BaseTransport`` with the retry policy above.""" 

98 

99 def __init__( 

100 self, 

101 *, 

102 transport: httpx.BaseTransport | None = None, 

103 retry: RetryConfig = DEFAULT_RETRY_CONFIG, 

104 ) -> None: 

105 self._transport = transport if transport is not None else httpx.HTTPTransport() 

106 self._retry = retry 

107 

108 def handle_request(self, request: httpx.Request) -> httpx.Response: 

109 # A bounded for-loop, not `while True` plus a hand-incremented 

110 # counter: `_should_retry` already returns False once `attempt` 

111 # reaches `max_retries`, so the last iteration always returns or 

112 # raises from inside the loop -- but driving termination off 

113 # `range()` too means a broken attempt counter can never turn this 

114 # into an infinite retry loop, only a wrong retry *count*. 

115 # The exact upper bound only has to be "enough" -- _should_retry's 

116 # own attempt >= max_retries check is what actually ends the loop 

117 # on the right iteration, so range()'s bound is unreachable as a 

118 # binding constraint whenever _should_retry itself is correct. 

119 for attempt in range(self._retry.max_retries + 1): # pragma: no mutate 

120 try: 

121 response = self._transport.handle_request(request) 

122 except httpx.TransportError as exc: 

123 if not _should_retry(attempt, self._retry, None, exc): 

124 raise 

125 time.sleep(_retry_delay(attempt, self._retry, None)) 

126 continue 

127 

128 if not _should_retry(attempt, self._retry, response, None): 

129 return response 

130 

131 response.close() 

132 time.sleep(_retry_delay(attempt, self._retry, response)) 

133 

134 # Every loop iteration above returns or raises before falling 

135 # through, so this is unreachable -- kept only so mypy sees an 

136 # explicit return on every path. 

137 raise AssertionError("unreachable") # pragma: no cover, no mutate 

138 

139 def close(self) -> None: 

140 self._transport.close() 

141 

142 

143class AsyncRetryTransport(httpx.AsyncBaseTransport): 

144 """Wraps an async ``httpx.AsyncBaseTransport`` with the retry policy above.""" 

145 

146 def __init__( 

147 self, 

148 *, 

149 transport: httpx.AsyncBaseTransport | None = None, 

150 retry: RetryConfig = DEFAULT_RETRY_CONFIG, 

151 ) -> None: 

152 self._transport = transport if transport is not None else httpx.AsyncHTTPTransport() 

153 self._retry = retry 

154 

155 async def handle_async_request(self, request: httpx.Request) -> httpx.Response: 

156 # See RetryTransport.handle_request's comments: bounded by 

157 # range(), not a hand-incremented counter under `while True`, and 

158 # the exact bound is unreachable as a binding constraint whenever 

159 # _should_retry itself is correct. 

160 for attempt in range(self._retry.max_retries + 1): # pragma: no mutate 

161 try: 

162 response = await self._transport.handle_async_request(request) 

163 except httpx.TransportError as exc: 

164 if not _should_retry(attempt, self._retry, None, exc): 

165 raise 

166 await asyncio.sleep(_retry_delay(attempt, self._retry, None)) 

167 continue 

168 

169 if not _should_retry(attempt, self._retry, response, None): 

170 return response 

171 

172 await response.aclose() 

173 await asyncio.sleep(_retry_delay(attempt, self._retry, response)) 

174 

175 raise AssertionError("unreachable") # pragma: no cover, no mutate 

176 

177 async def aclose(self) -> None: 

178 await self._transport.aclose()