fix(cdk): sync injected api_budget onto LimiterSession and guard with APIBudget isinstance

Co-Authored-By: bot_apk <apk@cognition.ai>

Author: devin-ai-integration[bot]

airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py

@@ -1884,10 +1884,14 @@ def create_custom_component(self, model: Any, config: Config, **kwargs: Any) ->
 
     @staticmethod
     def _sync_injected_api_budget_with_http_client(custom_requester: HttpRequester) -> None:
-        """
-        Custom requesters can replace `_http_client` in `__post_init__` without forwarding `api_budget`.
-        If the factory injected a manifest-level budget and the replacement client kept the default empty
-        budget, sync the active client back to the requester's injected budget.
+        """Align an injected `api_budget` with the active `HttpClient` on custom requesters.
+
+        Custom requesters can replace `_http_client` in `__post_init__` without forwarding
+        `api_budget`. If the factory injected a manifest-level budget and the replacement
+        client kept the default empty `APIBudget`, point both the client and its underlying
+        `LimiterSession`/`CachedLimiterSession` at the injected budget so rate-limiting is
+        actually enforced at request time. Non-`APIBudget` implementations (custom
+        `AbstractAPIBudget` subclasses) are left untouched.
         """
         http_client = getattr(custom_requester, "_http_client", None)
         http_client_api_budget = getattr(http_client, "_api_budget", None)
@@ -1901,8 +1905,17 @@ def _sync_injected_api_budget_with_http_client(custom_requester: HttpRequester)
         ):
             return
 
-        if len(getattr(http_client_api_budget, "_policies", [])) == 0:
+        if (
+            isinstance(http_client_api_budget, APIBudget)
+            and len(http_client_api_budget._policies) == 0
+        ):
             http_client._api_budget = injected_api_budget
+            http_client_session = getattr(http_client, "_session", None)
+            if (
+                http_client_session is not None
+                and getattr(http_client_session, "_api_budget", None) is http_client_api_budget
+            ):
+                http_client_session._api_budget = injected_api_budget
 
     @staticmethod
     def _get_class_from_fully_qualified_class_name(

unit_tests/sources/declarative/parsers/test_model_to_component_factory.py

@@ -4364,8 +4364,11 @@ def test_api_budget_propagated_to_custom_requester_subclass_of_http_requester():
     assert len(custom_requester.api_budget._policies) == 1
     policy = custom_requester.api_budget._policies[0]
     assert isinstance(policy, MovingWindowCallRatePolicy)
-    # Also verify the underlying HttpClient received the same budget
+    # Verify the underlying HttpClient AND its LimiterSession both received the same
+    # budget: rate-limiting is enforced on the session at send() time, so asserting only
+    # on the client field is insufficient to prove the policies are actually active.
     assert custom_requester._http_client._api_budget is custom_requester.api_budget
+    assert custom_requester._http_client._session._api_budget is custom_requester.api_budget
 
 
 def test_api_budget_propagated_to_custom_requester_that_replaces_http_client():
@@ -4414,6 +4417,10 @@ def test_api_budget_propagated_to_custom_requester_that_replaces_http_client():
     assert isinstance(custom_requester, HttpRequester)
     assert custom_requester.api_budget is not None
     assert custom_requester._http_client._api_budget is custom_requester.api_budget
+    # The LimiterSession holds its own reference to the budget (captured at client
+    # construction time) and is what actually enforces rate limits on send(). Assert
+    # it was synced too, otherwise the injected budget is effectively inert.
+    assert custom_requester._http_client._session._api_budget is custom_requester.api_budget
 
 
 def test_api_budget_not_overwriting_non_empty_budget_on_replaced_http_client():
@@ -4458,6 +4465,13 @@ def test_api_budget_not_overwriting_non_empty_budget_on_replaced_http_client():
     assert custom_requester.api_budget is not None
     assert custom_requester._http_client._api_budget is not custom_requester.api_budget
     assert len(custom_requester._http_client._api_budget._policies) == 1
+    # The client's own budget must remain wired into its LimiterSession as well, so
+    # the sync step never silently swaps an intentionally-installed budget out from
+    # under the active session.
+    assert (
+        custom_requester._http_client._session._api_budget
+        is custom_requester._http_client._api_budget
+    )
 
 
 def test_api_budget_not_propagated_to_non_http_requester_custom_components():