fix(call-rate): honor RateLimit-Remaining when a reset header is present

airbyte_cdk/sources/declarative/declarative_component_schema.yaml

@@ -2014,6 +2014,11 @@ definitions:
         type: array
         items:
           "$ref": "#/definitions/HttpRequestRegexMatcher"
+      max_header_driven_wait:
+        title: Maximum Header-Driven Wait
+        description: Maximum wait a rate-limit-header-driven update may induce. Rates with longer windows are not adjusted from headers. Defaults to PT10M.
+        type: string
+        examples: ["PT10M", "PT1M"]
     additionalProperties: true
   UnlimitedCallRatePolicy:
     title: Unlimited Call Rate Policy

airbyte_cdk/sources/declarative/models/declarative_component_schema.py

@@ -2217,6 +2217,12 @@ class Config:
         description="List of matchers that define which requests this policy applies to.",
         title="Matchers",
     )
+    max_header_driven_wait: Optional[str] = Field(
+        None,
+        description="Maximum wait a rate-limit-header-driven update may induce. Rates with longer windows are not adjusted from headers. Defaults to PT10M.",
+        examples=["PT10M", "PT1M"],
+        title="Maximum Header-Driven Wait",
+    )
 
 
 class UnlimitedCallRatePolicy(BaseModel):

airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py

@@ -4603,6 +4603,11 @@ def create_moving_window_call_rate_policy(
         return MovingWindowCallRatePolicy(
             rates=rates,
             matchers=matchers,
+            **(
+                {"max_header_driven_wait": parse_duration(model.max_header_driven_wait)}
+                if model.max_header_driven_wait
+                else {}
+            ),
         )
 
     def create_unlimited_call_rate_policy(

airbyte_cdk/sources/streams/call_rate.py

@@ -15,7 +15,7 @@
 
 import requests
 import requests_cache
-from pyrate_limiter import InMemoryBucket, Limiter, RateItem, TimeClock
+from pyrate_limiter import InMemoryBucket, Limiter, RateItem, TimeClock, binary_search
 from pyrate_limiter import Rate as PyRateRate
 from pyrate_limiter.exceptions import BucketFullException
 
@@ -429,19 +429,35 @@ class MovingWindowCallRatePolicy(BaseCallRatePolicy):
     This strategy requires saving of timestamps of all requests within a window.
     """
 
-    def __init__(self, rates: list[Rate], matchers: list[RequestMatcher]):
+    # Header-driven updates must not induce waits longer than this; sources heartbeat well above it.
+    # Policies with longer windows are left unchanged rather than parking a worker for the whole window.
+    DEFAULT_MAX_HEADER_DRIVEN_WAIT = timedelta(minutes=10)
+
+    def __init__(
+        self,
+        rates: list[Rate],
+        matchers: list[RequestMatcher],
+        max_header_driven_wait: timedelta = DEFAULT_MAX_HEADER_DRIVEN_WAIT,
+    ):
         """Constructor
 
         :param rates: list of rates, the order is important and must be ascending
         :param matchers:
+        :param max_header_driven_wait: maximum wait a header-driven update may induce. Rates whose
+            windows exceed this value are excluded from header-driven updates, so a policy with no
+            rate inside the bound is never adjusted from response headers. Defaults to 10 minutes,
+            chosen to stay well inside the platform's source heartbeat.
         """
         if not rates:
             raise ValueError("The list of rates can not be empty")
+        if max_header_driven_wait <= timedelta(0):
+            raise ValueError("max_header_driven_wait must be positive")
         pyrate_rates = [
             PyRateRate(limit=rate.limit, interval=int(rate.interval.total_seconds() * 1000))
             for rate in rates
         ]
         self._bucket = InMemoryBucket(pyrate_rates)
+        self._max_header_driven_wait_ms = int(max_header_driven_wait.total_seconds() * 1000)
         # Limiter will create the background task that clears old requests in the bucket
         self._limiter = Limiter(self._bucket)
         super().__init__(matchers=matchers)
@@ -478,22 +494,64 @@ def update(
     ) -> None:
         """Adjust call bucket to reflect the state of the API server
 
-        :param available_calls:
-        :param call_reset_ts:
+        The bucket is filled with dummy calls until what it still allows matches what the API
+        reports as available. Updates only ever lower the local allowance: when the API reports
+        more available calls than the configured rates allow, the configured rates win.
+
+        `call_reset_ts` is not used. A moving window has no reset point, so the only actionable
+        part of the API feedback is the number of calls left; the window length stays the one
+        the rates were configured with.
+
+        When several rates are configured, the update applies to the most constraining one. A
+        header describing a coarser window only starts to bite near the end of that window;
+        mapping headers to a specific rate is deliberately out of scope.
+
+        The subtraction below compares `_calls_left`, which counts bucket entries (weight units;
+        `put` stores `weight` copies), with `available_calls`, which counts requests. They
+        coincide for unweighted policies but diverge when weighted matchers are used with a
+        remaining header.
+
+        :param available_calls: number of calls the API reports as still available
+        :param call_reset_ts: unused, see above
         :return:
         """
-        if (
-            available_calls is not None and call_reset_ts is None
-        ):  # we do our best to sync buckets with API
-            if available_calls == 0:
-                with self._limiter.lock:
-                    items_to_add = self._bucket.count() < self._bucket.rates[0].limit
-                    if items_to_add > 0:
-                        now: int = TimeClock().now()  # type: ignore[no-untyped-call]
-                        self._bucket.put(RateItem(name="dummy", timestamp=now, weight=items_to_add))
-        # TODO: add support if needed, it might be that it is not possible to make a good solution for this case
-        # if available_calls is not None and call_reset_ts is not None:
-        #     ts = call_reset_ts.timestamp()
+        if available_calls is None:
+            return
+
+        available_calls = max(0, available_calls)
+        with self._limiter.lock:
+            now: int = TimeClock().now()  # type: ignore[no-untyped-call]
+            calls_left = self._calls_left(now)
+            if calls_left is None:
+                return
+
+            items_to_add = calls_left - available_calls
+            if items_to_add > 0:
+                logger.debug(
+                    "got rate limit update from api, adjusting available calls from %s to %s",
+                    calls_left,
+                    available_calls,
+                )
+                if not self._bucket.put(RateItem(name="dummy", timestamp=now, weight=items_to_add)):
+                    logger.warning(
+                        "could not adjust available calls: calls_left=%s, available_calls=%s, "
+                        "rejected_weight=%s",
+                        calls_left,
+                        available_calls,
+                        items_to_add,
+                    )
+
+    def _calls_left(self, now: int) -> Optional[int]:
+        """Number of calls the bucket still allows, i.e. the most constraining of all rates."""
+        items = self._bucket.items
+        calls_left = []
+        for rate in self._bucket.rates:
+            if rate.interval > self._max_header_driven_wait_ms:
+                continue
+            lower_bound_idx = binary_search(items, now - rate.interval)
+            calls_used = len(items) - lower_bound_idx if lower_bound_idx >= 0 else 0
+            calls_left.append(rate.limit - calls_used)
+        return min(calls_left) if calls_left else None
 
     def __str__(self) -> str:
         """Return a human-friendly description of the moving window rate policy for logging purposes."""

unit_tests/sources/declarative/parsers/test_model_to_component_factory.py

@@ -4879,6 +4879,7 @@ def test_api_budget():
                             "interval": "PT0.1S",  # 0.1 seconds
                         }
                     ],
+                    "max_header_driven_wait": "PT2M",
                     "matchers": [
                         {
                             "type": "HttpRequestRegexMatcher",
@@ -4938,6 +4939,7 @@ def test_api_budget():
     # The 0.1s from 'PT0.1S' is stored in ms by PyRateLimiter internally
     # but here just check that the limit and interval exist
     assert policy._bucket.rates[0].interval == 100  # 100 ms
+    assert policy._max_header_driven_wait_ms == 120_000
 
 
 def test_api_budget_passed_to_custom_requester():
@@ -4981,6 +4983,9 @@ def test_api_budget_passed_to_custom_requester():
     assert isinstance(custom_requester.api_budget, HttpAPIBudget)
     assert custom_requester._http_client._api_budget is custom_requester.api_budget
     assert len(custom_requester._http_client._api_budget._policies) == 1
+    policy = custom_requester.api_budget._policies[0]
+    assert isinstance(policy, MovingWindowCallRatePolicy)
+    assert policy._max_header_driven_wait_ms == 600_000
 
 
 def test_api_budget_does_not_override_custom_requester_default_value():

unit_tests/sources/streams/test_call_rate.py

@@ -16,6 +16,7 @@
     APIBudget,
     CallRateLimitHit,
     FixedWindowCallRatePolicy,
+    HttpAPIBudget,
     HttpRequestMatcher,
     HttpRequestRegexMatcher,
     MovingWindowCallRatePolicy,
@@ -256,6 +257,15 @@ def test_no_rates(self):
         with pytest.raises(ValueError, match="The list of rates can not be empty"):
             MovingWindowCallRatePolicy(rates=[], matchers=[])
 
+    @pytest.mark.parametrize("max_header_driven_wait", [timedelta(0), timedelta(minutes=-1)])
+    def test_invalid_max_header_driven_wait(self, max_header_driven_wait):
+        with pytest.raises(ValueError, match="max_header_driven_wait must be positive"):
+            MovingWindowCallRatePolicy(
+                rates=[Rate(10, timedelta(minutes=1))],
+                matchers=[],
+                max_header_driven_wait=max_header_driven_wait,
+            )
+
     def test_limit_rate(self):
         """try_acquire must respect configured call rate and throw CallRateLimitHit when hit the limit."""
         policy = MovingWindowCallRatePolicy(rates=[Rate(10, timedelta(minutes=1))], matchers=[])
@@ -305,6 +315,170 @@ def test_multiple_limit_rates(self):
         assert excinfo.value.time_to_wait.total_seconds() == pytest.approx(3600, 0.1)
         assert str(excinfo.value) == "Bucket for item=call with Rate limit=2/1.0h is already full"
 
+    def test_update_available_calls_with_reset_ts(self):
+        policy = MovingWindowCallRatePolicy(rates=[Rate(10, timedelta(minutes=1))], matchers=[])
+
+        policy.update(available_calls=2, call_reset_ts=datetime.now())
+
+        policy.try_acquire("call", weight=1)
+        policy.try_acquire("call", weight=1)
+        with pytest.raises(CallRateLimitHit):
+            policy.try_acquire("call", weight=1)
+
+    def test_update_only_lowers_allowance(self):
+        policy = MovingWindowCallRatePolicy(rates=[Rate(10, timedelta(minutes=1))], matchers=[])
+
+        policy.update(available_calls=50, call_reset_ts=datetime.now())
+
+        for _ in range(10):
+            policy.try_acquire("call", weight=1)
+        with pytest.raises(CallRateLimitHit):
+            policy.try_acquire("call", weight=1)
+
+    def test_update_is_noop_without_available_calls(self):
+        policy = MovingWindowCallRatePolicy(rates=[Rate(10, timedelta(minutes=1))], matchers=[])
+
+        policy.update(available_calls=None, call_reset_ts=datetime.now())
+        policy.update(available_calls=None, call_reset_ts=None)
+
+        for _ in range(10):
+            policy.try_acquire("call", weight=1)
+
+    def test_update_respects_the_most_constraining_rate(self):
+        policy = MovingWindowCallRatePolicy(
+            rates=[
+                Rate(10, timedelta(seconds=1)),
+                Rate(5, timedelta(minutes=1)),
+            ],
+            matchers=[],
+        )
+
+        policy.update(available_calls=1, call_reset_ts=None)
+
+        policy.try_acquire("call", weight=1)
+        with pytest.raises(CallRateLimitHit) as exc:
+            policy.try_acquire("call", weight=1)
+        assert exc.value.rate == "limit=5/1.0m"
+        assert exc.value.time_to_wait.total_seconds() == pytest.approx(60, 0.1)
+
+    def test_update_available_calls_zero_fills_bucket(self):
+        policy = MovingWindowCallRatePolicy(rates=[Rate(10, timedelta(minutes=1))], matchers=[])
+
+        policy.update(available_calls=0, call_reset_ts=None)
+
+        with pytest.raises(CallRateLimitHit):
+            policy.try_acquire("call", weight=1)
+
+    def test_update_ignores_rates_over_header_wait_cap(self):
+        policy = MovingWindowCallRatePolicy(rates=[Rate(100, timedelta(minutes=15))], matchers=[])
+
+        policy.update(available_calls=0, call_reset_ts=None)
+
+        for _ in range(100):
+            policy.try_acquire("call", weight=1)
+
+    def test_update_uses_configured_header_wait_cap(self):
+        policy = MovingWindowCallRatePolicy(
+            rates=[Rate(100, timedelta(minutes=15))],
+            matchers=[],
+            max_header_driven_wait=timedelta(minutes=20),
+        )
+
+        policy.update(available_calls=0, call_reset_ts=None)
+
+        with pytest.raises(CallRateLimitHit):
+            policy.try_acquire("call", weight=1)
+
+    def test_update_uses_tightened_header_wait_cap(self):
+        default_policy = MovingWindowCallRatePolicy(
+            rates=[Rate(100, timedelta(minutes=5))], matchers=[]
+        )
+        tightened_policy = MovingWindowCallRatePolicy(
+            rates=[Rate(100, timedelta(minutes=5))],
+            matchers=[],
+            max_header_driven_wait=timedelta(minutes=1),
+        )
+
+        default_policy.update(available_calls=0, call_reset_ts=None)
+        tightened_policy.update(available_calls=0, call_reset_ts=None)
+
+        with pytest.raises(CallRateLimitHit):
+            default_policy.try_acquire("call", weight=1)
+        for _ in range(100):
+            tightened_policy.try_acquire("call", weight=1)
+
+    def test_update_caps_to_eligible_rate(self):
+        policy = MovingWindowCallRatePolicy(
+            rates=[
+                Rate(10, timedelta(minutes=1)),
+                Rate(100, timedelta(minutes=15)),
+            ],
+            matchers=[],
+        )
+
+        policy.update(available_calls=0, call_reset_ts=None)
+
+        with pytest.raises(CallRateLimitHit) as exc:
+            policy.try_acquire("call", weight=1)
+        assert exc.value.time_to_wait.total_seconds() <= 600
+
+    def test_update_clamps_negative_available_calls(self):
+        policy = MovingWindowCallRatePolicy(rates=[Rate(10, timedelta(minutes=1))], matchers=[])
+
+        policy.update(available_calls=-1, call_reset_ts=None)
+
+        with pytest.raises(CallRateLimitHit):
+            policy.try_acquire("call", weight=1)
+
+
+class TestHttpAPIBudget:
+    def test_update_from_response(self, mocker):
+        policy = MovingWindowCallRatePolicy(rates=[Rate(10, timedelta(minutes=1))], matchers=[])
+        budget = HttpAPIBudget(policies=[policy])
+        request = Request("GET", "https://example.com")
+        response = mocker.Mock(spec=requests.Response)
+        response.headers = requests.structures.CaseInsensitiveDict(
+            {
+                "RateLimit-Remaining": "1",
+                "RateLimit-Reset": "60",
+                "RateLimit-Limit": "60",
+            }
+        )
+        response.status_code = 200
+
+        budget.update_from_response(request, response)
+
+        budget.acquire_call(request, block=False)
+        with pytest.raises(CallRateLimitHit):
+            budget.acquire_call(request, block=False)
+
+    def test_update_from_429_response(self, mocker):
+        policy = MovingWindowCallRatePolicy(rates=[Rate(10, timedelta(minutes=1))], matchers=[])
+        budget = HttpAPIBudget(policies=[policy])
+        request = Request("GET", "https://example.com")
+        response = mocker.Mock(spec=requests.Response)
+        response.headers = requests.structures.CaseInsensitiveDict()
+        response.status_code = 429
+
+        budget.update_from_response(request, response)
+
+        with pytest.raises(CallRateLimitHit) as exc:
+            budget.acquire_call(request, block=False)
+        assert exc.value.time_to_wait.total_seconds() <= 600
+
+    def test_update_from_429_response_ignores_over_cap_policy(self, mocker):
+        policy = MovingWindowCallRatePolicy(rates=[Rate(100, timedelta(minutes=15))], matchers=[])
+        budget = HttpAPIBudget(policies=[policy])
+        request = Request("GET", "https://example.com")
+        response = mocker.Mock(spec=requests.Response)
+        response.headers = requests.structures.CaseInsensitiveDict()
+        response.status_code = 429
+
+        budget.update_from_response(request, response)
+
+        for _ in range(100):
+            budget.acquire_call(request, block=False)
+
 
 class TestHttpStreamIntegration:
     def test_without_cache(self, mocker, requests_mock):