Exponential Backoff & Fail-Stop on 401 (Invalid API Token) (#342)

Author: IgorChvyrov-sm

docs/authentication-retry-policy.md

@@ -14,14 +14,33 @@
 from conductor.asyncio_client.http.api_response import ApiResponse
 from conductor.asyncio_client.http.api_response import T as ApiResponseT
 from conductor.asyncio_client.http.exceptions import ApiException
+from conductor.client.exceptions.auth_401_policy import Auth401Policy, Auth401Handler
 
 logger = logging.getLogger(Configuration.get_logging_formatted_name(__name__))
 
 
 class ApiClientAdapter(ApiClient):
-    def __init__(self, *args, **kwargs):
+    def __init__(
+        self, configuration=None, header_name=None, header_value=None, cookie=None
+    ):
         self._token_lock = asyncio.Lock()
-        super().__init__(*args, **kwargs)
+        self.configuration = configuration or Configuration()
+
+        self.rest_client = rest.RESTClientObject(self.configuration)
+        self.default_headers = {}
+        if header_name is not None:
+            self.default_headers[header_name] = header_value
+        self.cookie = cookie
+
+        # Initialize 401 policy handler
+        auth_401_policy = Auth401Policy(
+            max_attempts=self.configuration.auth_401_max_attempts,
+            base_delay_ms=self.configuration.auth_401_base_delay_ms,
+            max_delay_ms=self.configuration.auth_401_max_delay_ms,
+            jitter_percent=self.configuration.auth_401_jitter_percent,
+            stop_behavior=self.configuration.auth_401_stop_behavior,
+        )
+        self.auth_401_handler = Auth401Handler(auth_401_policy)
 
     async def call_api(
         self,
@@ -46,7 +65,10 @@ async def call_api(
 
         try:
             logger.debug(
-                "HTTP request method: %s; url: %s; header_params: %s", method, url, header_params
+                "HTTP request method: %s; url: %s; header_params: %s",
+                method,
+                url,
+                header_params,
             )
             response_data = await self.rest_client.request(
                 method,
@@ -56,41 +78,133 @@ async def call_api(
                 post_params=post_params,
                 _request_timeout=_request_timeout,
             )
-            if (
+
+            # Handle 401 retries with policy-based logic
+            resource_path = url.replace(self.configuration.host, "")
+
+            # Loop to handle multiple 401 retries
+            while (
                 response_data.status == 401  # noqa: PLR2004 (Unauthorized status code)
                 and url != self.configuration.host + "/token"
             ):
-                logger.warning(
-                    "HTTP response from: %s; status code: 401 - obtaining new token", url
-                )
-                async with self._token_lock:
-                    # The lock is intentionally broad (covers the whole block including the token state)
-                    # to avoid race conditions: without it, other coroutines could mis-evaluate
-                    # token state during a context switch and trigger redundant refreshes
-                    token_expired = (
-                        self.configuration.token_update_time > 0
-                        and time.time()
-                        >= self.configuration.token_update_time
-                        + self.configuration.auth_token_ttl_sec
+                # Check if this is an auth-dependent call that should trigger 401 policy
+                if self.auth_401_handler.policy.is_auth_dependent_call(
+                    resource_path, method
+                ):
+                    # Handle 401 with policy (exponential backoff, max attempts, etc.)
+                    result = self.auth_401_handler.handle_401_error(
+                        resource_path=resource_path,
+                        method=method,
+                        status_code=401,
+                        error_code=None,
                     )
-                    invalid_token = not self.configuration._http_config.api_key.get("api_key")
 
-                    if invalid_token or token_expired:
-                        token = await self.refresh_authorization_token()
+                    if result["should_retry"]:
+                        # Apply exponential backoff delay
+                        if result["delay_seconds"] > 0:
+                            logger.info(
+                                "401 error on %s %s - waiting %.2fs before retry (attempt %d/%d)",
+                                method,
+                                url,
+                                result["delay_seconds"],
+                                result["attempt_count"],
+                                result["max_attempts"],
+                            )
+                            await asyncio.sleep(result["delay_seconds"])
+
+                        # Try to refresh token and retry
+                        async with self._token_lock:
+                            # Check if token was already refreshed by another coroutine
+                            # to avoid race condition where multiple concurrent 401s
+                            # trigger redundant token refreshes
+                            token_expired = (
+                                self.configuration.token_update_time > 0
+                                and time.time()
+                                >= self.configuration.token_update_time
+                                + self.configuration.auth_token_ttl_sec
+                            )
+                            invalid_token = (
+                                not self.configuration._http_config.api_key.get(
+                                    "api_key"
+                                )
+                            )
+
+                            if invalid_token or token_expired:
+                                token = await self.refresh_authorization_token()
+                            else:
+                                token = self.configuration._http_config.api_key[
+                                    "api_key"
+                                ]
+                            if header_params is None:
+                                header_params = {}
+                            header_params["X-Authorization"] = token
+
+                        # Make the retry request outside the lock to avoid blocking other coroutines
+                        response_data = await self.rest_client.request(
+                            method,
+                            url,
+                            headers=header_params,
+                            body=body,
+                            post_params=post_params,
+                            _request_timeout=_request_timeout,
+                        )
                     else:
-                        token = self.configuration._http_config.api_key["api_key"]
-                    header_params["X-Authorization"] = token
-                    response_data = await self.rest_client.request(
-                        method,
+                        # Max attempts reached - log error and break
+                        logger.error(
+                            "401 error on %s %s - max attempts (%d) reached, stopping worker",
+                            method,
+                            url,
+                            result["max_attempts"],
+                        )
+                        break
+                else:
+                    # Non-auth-dependent call with 401 - use original behavior (single retry)
+                    logger.warning(
+                        "HTTP response from: %s; status code: 401 - obtaining new token",
                         url,
-                        headers=header_params,
-                        body=body,
-                        post_params=post_params,
-                        _request_timeout=_request_timeout,
                     )
+                    async with self._token_lock:
+                        # The lock is intentionally broad (covers the whole block including the token state)
+                        # to avoid race conditions: without it, other coroutines could mis-evaluate
+                        # token state during a context switch and trigger redundant refreshes
+                        token_expired = (
+                            self.configuration.token_update_time > 0
+                            and time.time()
+                            >= self.configuration.token_update_time
+                            + self.configuration.auth_token_ttl_sec
+                        )
+                        invalid_token = not self.configuration._http_config.api_key.get(
+                            "api_key"
+                        )
+
+                        if invalid_token or token_expired:
+                            token = await self.refresh_authorization_token()
+                        else:
+                            token = self.configuration._http_config.api_key["api_key"]
+                        if header_params is None:
+                            header_params = {}
+                        header_params["X-Authorization"] = token
+                        response_data = await self.rest_client.request(
+                            method,
+                            url,
+                            headers=header_params,
+                            body=body,
+                            post_params=post_params,
+                            _request_timeout=_request_timeout,
+                        )
+                    # Break after single retry for non-auth-dependent calls
+                    break
+
+            # Record successful call to reset 401 attempt counters
+            if response_data.status != 401:
+                self.auth_401_handler.record_successful_call(resource_path)
+
         except ApiException as e:
             logger.error(
-                "HTTP request failed url: %s status: %s; reason: %s", url, e.status, e.reason
+                "HTTP request failed url: %s status: %s; reason: %s",
+                url,
+                e.status,
+                e.reason,
             )
             raise e
 
@@ -117,7 +231,9 @@ def response_deserialize(
             and 100 <= response_data.status <= 599  # noqa: PLR2004
         ):
             # if not found, look for '1XX', '2XX', etc.
-            response_type = response_types_map.get(str(response_data.status)[0] + "XX", None)
+            response_type = response_types_map.get(
+                str(response_data.status)[0] + "XX", None
+            )
 
         # deserialize response data
         response_text = None
@@ -134,10 +250,14 @@ def response_deserialize(
                     match = re.search(r"charset=([a-zA-Z\-\d]+)[\s;]?", content_type)
                 encoding = match.group(1) if match else "utf-8"
                 response_text = response_data.data.decode(encoding)
-                return_data = self.deserialize(response_text, response_type, content_type)
+                return_data = self.deserialize(
+                    response_text, response_type, content_type
+                )
         finally:
             if not 200 <= response_data.status <= 299:  #  noqa: PLR2004
-                logger.error("Unexpected response status code: %s", response_data.status)
+                logger.error(
+                    "Unexpected response status code: %s", response_data.status
+                )
                 raise ApiException.from_response(
                     http_resp=response_data,
                     body=response_text,

src/conductor/asyncio_client/adapters/api_client_adapter.py

@@ -64,6 +64,13 @@ async def run(self) -> None:
         )
 
         while True:
+            # Check if worker should stop due to 401 policy
+            if (hasattr(self.task_client, 'api_client') and 
+                hasattr(self.task_client.api_client, 'auth_401_handler') and 
+                hasattr(self.task_client.api_client.auth_401_handler, 'is_worker_stopped') and 
+                self.task_client.api_client.auth_401_handler.is_worker_stopped()):
+                logger.error("Worker stopped due to persistent 401 authentication failures")
+                break
             await self.run_once()
 
     async def run_once(self) -> None:

src/conductor/asyncio_client/automator/task_runner.py

@@ -33,6 +33,15 @@ class Configuration:
     CONDUCTOR_WORKER_<TASK_TYPE>_POLLING_INTERVAL: Task-specific polling interval
     CONDUCTOR_WORKER_<TASK_TYPE>_DOMAIN: Task-specific domain
 
+    Authentication Retry Policy (401 Handling):
+    -------------------------------------------
+    CONDUCTOR_AUTH_401_MAX_ATTEMPTS: Maximum retry attempts per endpoint (default: 6)
+    CONDUCTOR_AUTH_401_BASE_DELAY_MS: Base delay in milliseconds (default: 1000.0)
+    CONDUCTOR_AUTH_401_MAX_DELAY_MS: Maximum delay cap in milliseconds (default: 60000.0)
+    CONDUCTOR_AUTH_401_JITTER_PERCENT: Random jitter percentage 0.0-1.0 (default: 0.2)
+    CONDUCTOR_AUTH_401_STOP_BEHAVIOR: Behavior after max attempts: 'stop_worker' or 'continue' (default: 'stop_worker')
+
+    
     Example:
     --------
     ```python
@@ -82,6 +91,12 @@ def __init__(
         verify_ssl: Optional[bool] = None,
         proxy: Optional[str] = None,
         proxy_headers: Optional[Dict[str, str]] = None,
+        # 401-specific configuration
+        auth_401_max_attempts: Optional[int] = None,
+        auth_401_base_delay_ms: Optional[float] = None,
+        auth_401_max_delay_ms: Optional[float] = None,
+        auth_401_jitter_percent: Optional[float] = None,
+        auth_401_stop_behavior: Optional[str] = None,
         **kwargs: Any,
     ):
         """
@@ -233,6 +248,23 @@ def __init__(
         # Orkes Conductor auth token properties
         self.token_update_time = 0
         self.auth_token_ttl_sec = auth_token_ttl_min * 60
+        
+        # 401-specific configuration
+        self.auth_401_max_attempts = auth_401_max_attempts or self._get_env_int(
+            "CONDUCTOR_AUTH_401_MAX_ATTEMPTS", 6
+        )
+        self.auth_401_base_delay_ms = auth_401_base_delay_ms or self._get_env_float(
+            "CONDUCTOR_AUTH_401_BASE_DELAY_MS", 1000.0
+        )
+        self.auth_401_max_delay_ms = auth_401_max_delay_ms or self._get_env_float(
+            "CONDUCTOR_AUTH_401_MAX_DELAY_MS", 60000.0
+        )
+        self.auth_401_jitter_percent = auth_401_jitter_percent or self._get_env_float(
+            "CONDUCTOR_AUTH_401_JITTER_PERCENT", 0.2
+        )
+        self.auth_401_stop_behavior = auth_401_stop_behavior or os.getenv(
+            "CONDUCTOR_AUTH_401_STOP_BEHAVIOR", "stop_worker"
+        )
 
     def _get_env_float(self, env_var: str, default: float) -> float:
         """Get float value from environment variable with default fallback."""