feat(ciba): handle the Retry-After header value upon slow_down error

Author: pmalouin

packages/auth0-ai/auth0_ai/authorizers/ciba/ciba_authorizer_base.py

@@ -154,6 +154,31 @@ async def _start(self, authorize_params) -> CIBAAuthorizationRequest:
             else:
                 raise
 
+    def _extract_retry_after_header(self, error: Auth0Error) -> Optional[int]:
+        """
+        Extract the Retry-After header value from an Auth0Error.
+        
+        Args:
+            error: The Auth0Error object that may contain HTTP headers
+            
+        Returns:
+            The retry-after value in seconds as an integer, or None if not present
+        """
+
+        if not hasattr(error, 'headers') or not error.headers:
+            return None
+            
+        retry_after = error.headers.get('retry-after') or error.headers.get('Retry-After')
+        
+        if retry_after is None:
+            return None
+            
+        try:
+            return int(retry_after)
+        except (ValueError, TypeError):
+            # If the retry-after value is not a valid integer, return None
+            return None
+
     def _get_credentials_internal(self, auth_request: CIBAAuthorizationRequest) -> TokenResponse | None:
         try:
             # Calculate elapsed time in seconds
@@ -180,7 +205,8 @@ def _get_credentials_internal(self, auth_request: CIBAAuthorizationRequest) -> T
                 raise AuthorizationPendingInterrupt(e.message, auth_request)
 
             if e.error_code == "slow_down":
-                raise AuthorizationPollingInterrupt(e.message, auth_request)
+                retry_after = self._extract_retry_after_header(e)
+                raise AuthorizationPollingInterrupt(e.message, auth_request, retry_after)
 
             if e.error_code == "invalid_grant":
                 raise InvalidGrantInterrupt(e.message, auth_request)
@@ -203,7 +229,7 @@ async def get_credentials_polling(self, auth_request: CIBAAuthorizationRequest)
             try:
                 credentials = self._get_credentials_internal(auth_request)
             except (AuthorizationPendingInterrupt, AuthorizationPollingInterrupt) as err:
-                await asyncio.sleep(err.request["interval"])
+                await asyncio.sleep(err.next_retry_interval())
             except Exception:
                 raise
 

packages/auth0-ai/auth0_ai/interrupts/ciba_interrupts.py

@@ -76,13 +76,24 @@ def __init__(self, message: str, request: CIBAAuthorizationRequest):
         CIBAInterrupt.__init__(self, message, AuthorizationPendingInterrupt.code)
         WithRequestData.__init__(self, request)
 
+    def next_retry_interval(self) -> int:
+        """Return the interval in seconds to wait before the next retry attempt."""
+        return self.request["interval"]
+
 
 class AuthorizationPollingInterrupt(CIBAInterrupt, WithRequestData):
     code: str = "CIBA_AUTHORIZATION_POLLING_ERROR"
 
-    def __init__(self, message: str, request: CIBAAuthorizationRequest):
+    def __init__(self, message: str, request: CIBAAuthorizationRequest, retry_after: int = None):
         Auth0Interrupt.__init__(self, message, AuthorizationPollingInterrupt.code)
         WithRequestData.__init__(self, request)
+        self.retry_after = retry_after
+
+    def next_retry_interval(self) -> int:
+        """Return the interval in seconds to wait before the next retry attempt."""
+        # Use the retry_after value from the HTTP header if available,
+        # otherwise fall back to the original interval from the auth request
+        return self.retry_after if self.retry_after is not None else self.request["interval"]
 
 
 class InvalidGrantInterrupt(CIBAInterrupt, WithRequestData):