diff --git a/src/zai/core/_http_client.py b/src/zai/core/_http_client.py index 080f8b2..706a50a 100644 --- a/src/zai/core/_http_client.py +++ b/src/zai/core/_http_client.py @@ -464,6 +464,26 @@ def _process_response_data( def _should_stream_response_body(self, request: httpx.Request) -> bool: return request.headers.get(RAW_RESPONSE_HEADER) == 'stream' # type: ignore[no-any-return] + def _business_error_code(self, response: httpx.Response) -> str | None: + try: + payload = response.json() + except httpx.ResponseNotRead: + try: + response.read() + payload = response.json() + except (httpx.HTTPError, ValueError): + return None + except ValueError: + return None + + if not isinstance(payload, Mapping): + return None + error = payload.get('error') + if not isinstance(error, Mapping): + return None + code = error.get('code') + return str(code) if code is not None else None + def _should_retry(self, response: httpx.Response) -> bool: # Note: this is not a standard header should_retry_header = response.headers.get('x-should-retry') @@ -488,6 +508,11 @@ def _should_retry(self, response: httpx.Response) -> bool: # Retry on rate limits. if response.status_code == 429: + # Business error 1113 requires an account recharge; retrying the same + # request cannot succeed without an external account-state change. + if self._business_error_code(response) == '1113': + log.debug('Not retrying HTTP 429 with non-retryable business error 1113') + return False log.debug('Retrying due to status code %i', response.status_code) return True diff --git a/tests/unit_tests/test_retry_policy.py b/tests/unit_tests/test_retry_policy.py new file mode 100644 index 0000000..9a252e1 --- /dev/null +++ b/tests/unit_tests/test_retry_policy.py @@ -0,0 +1,93 @@ +from __future__ import annotations + +import httpx +import pytest +from zai.core import APIReachLimitError +from zai.core._constants import ZAI_DEFAULT_TIMEOUT +from zai.core._http_client import HttpClient +from zai.core._streaming import StreamResponse + + +def make_client(http_client: httpx.Client | None = None, max_retries: int = 3) -> HttpClient: + return HttpClient( + version='test-version', + base_url='https://api.test.com', + _strict_response_validation=False, + max_retries=max_retries, + timeout=ZAI_DEFAULT_TIMEOUT, + custom_httpx_client=http_client, + ) + + +def test_http_429_arrears_is_not_retried() -> None: + attempts = 0 + + def handler(request: httpx.Request) -> httpx.Response: + nonlocal attempts + attempts += 1 + return httpx.Response( + 429, + json={'error': {'code': '1113', 'message': 'Account in arrears'}}, + request=request, + ) + + with httpx.Client(transport=httpx.MockTransport(handler)) as http_client: + client = make_client(http_client=http_client) + with pytest.raises(APIReachLimitError): + client.get('/resource', cast_type=dict) + + assert attempts == 1 + + +@pytest.mark.parametrize('code', ['1305', 1305]) +def test_http_429_rate_limit_remains_retryable(code: str | int) -> None: + client = make_client() + try: + response = httpx.Response(429, json={'error': {'code': code, 'message': 'Rate limited'}}) + assert client._should_retry(response) is True + finally: + client.close() + + +def test_http_429_unknown_payload_remains_retryable() -> None: + client = make_client() + try: + response = httpx.Response(429, content=b'not-json') + assert client._should_retry(response) is True + finally: + client.close() + + +@pytest.mark.parametrize(('header', 'expected'), [('true', True), ('false', False)]) +def test_explicit_retry_header_takes_precedence_for_arrears(header: str, expected: bool) -> None: + client = make_client() + try: + response = httpx.Response( + 429, + headers={'x-should-retry': header}, + json={'error': {'code': '1113', 'message': 'Account in arrears'}}, + ) + assert client._should_retry(response) is expected + finally: + client.close() + + +def test_streamed_http_429_arrears_is_not_retried() -> None: + attempts = 0 + + def handler(request: httpx.Request) -> httpx.Response: + nonlocal attempts + attempts += 1 + return httpx.Response( + 429, + request=request, + stream=httpx.ByteStream(b'{"error":{"code":"1113","message":"Account in arrears"}}'), + ) + + with httpx.Client(transport=httpx.MockTransport(handler)) as http_client: + client = make_client(http_client=http_client) + with pytest.raises(APIReachLimitError) as exc_info: + client.get('/resource', cast_type=dict, stream=True, stream_cls=StreamResponse) + + assert attempts == 1 + assert exc_info.value.response.json()['error']['code'] == '1113'