From 4ac11e2c823d8d58049e227d4f92bc3084a2d902 Mon Sep 17 00:00:00 2001 From: Nishant Date: Tue, 29 Sep 2026 06:22:06 +0000 Subject: [PATCH] zulip: Optionally retry requests that hit RATE_LIMIT_HIT. HTTP 429 currently bubbles up to every application. Tim asked for this as an opt-in Client option so callers that cannot block keep the old behavior. When auto_retry_rate_limits=True, do_api_query sleeps for the server's retry-after (JSON field, then Retry-After header) plus a small jitter, then retries, up to 10 times. The cap and jitter avoid pinning a client to the server's allowed rate. Fixes #773. Signed-off-by: Nishant --- zulip/tests/test_rate_limit.py | 117 +++++++++++++++++++++++++++++++++ zulip/zulip/__init__.py | 45 +++++++++++++ 2 files changed, 162 insertions(+) create mode 100644 zulip/tests/test_rate_limit.py diff --git a/zulip/tests/test_rate_limit.py b/zulip/tests/test_rate_limit.py new file mode 100644 index 000000000..bfb7534dd --- /dev/null +++ b/zulip/tests/test_rate_limit.py @@ -0,0 +1,117 @@ +import json +from typing import Any, Dict, Optional +from unittest import TestCase +from unittest.mock import Mock, patch + +import zulip + + +def _make_client(**kwargs: Any) -> zulip.Client: + with patch.object( + zulip.Client, + "get_server_settings", + return_value={"zulip_version": "10.0.0", "zulip_feature_level": 0}, + ): + return zulip.Client( + email="bot@example.com", + api_key="x" * 32, + site="https://example.zulipchat.com", + **kwargs, + ) + + +def _response( + status_code: int, + payload: Optional[Dict[str, Any]] = None, + headers: Optional[Dict[str, str]] = None, +) -> Mock: + res = Mock() + res.status_code = status_code + res.headers = headers or {} + if payload is None: + res.json.side_effect = json.JSONDecodeError("Expecting value", "", 0) + else: + res.json.return_value = payload + return res + + +class TestParseRetryAfter(TestCase): + def test_json_retry_after_preferred(self) -> None: + res = _response( + 429, + { + "result": "error", + "code": "RATE_LIMIT_HIT", + "msg": "rate limited", + "retry-after": 2.5, + }, + {"Retry-After": "9"}, + ) + self.assertEqual(zulip._parse_retry_after_seconds(res), 2.5) # noqa: SLF001 + + def test_retry_after_header_fallback(self) -> None: + res = _response(429, {"result": "error", "code": "RATE_LIMIT_HIT", "msg": "rate limited"}) + res.headers = {"Retry-After": "3.25"} + self.assertEqual(zulip._parse_retry_after_seconds(res), 3.25) # noqa: SLF001 + + def test_invalid_values_default_to_one(self) -> None: + res = _response(429, None, {"Retry-After": "soon"}) + self.assertEqual(zulip._parse_retry_after_seconds(res), 1.0) # noqa: SLF001 + + def test_negative_values_default_to_one(self) -> None: + res = _response(429, {"retry-after": -4}) + self.assertEqual(zulip._parse_retry_after_seconds(res), 1.0) # noqa: SLF001 + + +class TestAutoRetryRateLimits(TestCase): + def test_default_does_not_retry_429(self) -> None: + client = _make_client() + client.session = Mock() + client.session.request.return_value = _response( + 429, + {"result": "error", "code": "RATE_LIMIT_HIT", "msg": "API usage exceeded rate limit"}, + ) + with patch("time.sleep") as sleep: + result = client.do_api_query({}, "v1/messages") + sleep.assert_not_called() + self.assertEqual(result["code"], "RATE_LIMIT_HIT") + self.assertEqual(client.session.request.call_count, 1) + + def test_retries_after_retry_after_then_succeeds(self) -> None: + client = _make_client(auto_retry_rate_limits=True) + client.session = Mock() + limited = _response( + 429, + { + "result": "error", + "code": "RATE_LIMIT_HIT", + "msg": "API usage exceeded rate limit", + "retry-after": 1.5, + }, + ) + ok = _response(200, {"result": "success", "id": 1}) + client.session.request.side_effect = [limited, ok] + with patch("time.sleep") as sleep, patch("random.uniform", return_value=0.1): + result = client.do_api_query({}, "v1/messages") + sleep.assert_called_once_with(1.6) + self.assertEqual(result["result"], "success") + self.assertEqual(client.session.request.call_count, 2) + + def test_stops_after_ten_rate_limit_retries(self) -> None: + client = _make_client(auto_retry_rate_limits=True) + client.session = Mock() + limited = _response( + 429, + { + "result": "error", + "code": "RATE_LIMIT_HIT", + "msg": "rate limited", + "retry-after": 0.01, + }, + ) + client.session.request.return_value = limited + with patch("time.sleep"), patch("random.uniform", return_value=0.0): + result = client.do_api_query({}, "v1/messages") + # 10 retries after the first 429 => 11 requests, then the error is returned. + self.assertEqual(client.session.request.call_count, 11) + self.assertEqual(result["code"], "RATE_LIMIT_HIT") diff --git a/zulip/zulip/__init__.py b/zulip/zulip/__init__.py index 6d725a10b..1479bdc6a 100644 --- a/zulip/zulip/__init__.py +++ b/zulip/zulip/__init__.py @@ -385,6 +385,31 @@ class TimeZoneMissingError(ZulipError): pass +def _parse_retry_after_seconds(res: requests.Response) -> float: + """Return how long to wait after a 429, from JSON or the Retry-After header. + + Zulip documents `retry-after` on RATE_LIMIT_HIT responses: + https://zulip.com/api/rest-error-handling#rate-limit-exceeded + """ + wait: Optional[float] = None + try: + payload = res.json() + if isinstance(payload, dict) and payload.get("retry-after") is not None: + wait = float(payload["retry-after"]) + except (TypeError, ValueError, json.JSONDecodeError): + wait = None + if wait is None: + header = res.headers.get("Retry-After") + if header is not None: + try: + wait = float(header) + except (TypeError, ValueError): + wait = None + if wait is None or wait < 0: + wait = 1.0 + return wait + + class Client: def __init__( self, @@ -393,6 +418,7 @@ def __init__( config_file: Optional[str] = None, verbose: bool = False, retry_on_errors: bool = True, + auto_retry_rate_limits: bool = False, site: Optional[str] = None, client: Optional[str] = None, cert_bundle: Optional[str] = None, @@ -491,6 +517,7 @@ def __init__( self.base_url += "/api" self.base_url += "/" self.retry_on_errors = retry_on_errors + self.auto_retry_rate_limits = auto_retry_rate_limits self.client_name = client if insecure: @@ -607,6 +634,7 @@ def do_api_query( "had_error_retry": False, "request": request, "failures": 0, + "rate_limit_retries": 0, } def error_retry(error_string: str) -> bool: @@ -655,6 +683,23 @@ def end_error_retry(succeeded: bool) -> None: self.has_connected = True + # Opt-in: sleep for retry-after (plus a little jitter) and + # retry RATE_LIMIT_HIT. Default off so callers that cannot + # block are unchanged. Cap retries so a tight limit cannot + # pin a client to the server's allowed rate forever. + if ( + res.status_code == 429 + and self.auto_retry_rate_limits + and query_state["rate_limit_retries"] < 10 + ): + wait = _parse_retry_after_seconds(res) + wait += random.uniform(0.05, 0.5) # noqa: S311 + if self.verbose: + print(f"Rate limit hit; retrying after {wait:.2f}s.") + query_state["rate_limit_retries"] += 1 + time.sleep(wait) + continue + # On 50x errors, try again after a short sleep if str(res.status_code).startswith("5") and error_retry( f" (server {res.status_code})"