Prevent race condition during JWT obtaining

src/conductor/asyncio_client/adapters/api_client_adapter.py

@@ -1,5 +1,7 @@
+import asyncio
 import json
 import logging
+import time
 
 from conductor.asyncio_client.adapters.models import GenerateTokenRequest
 from conductor.asyncio_client.http import rest
@@ -10,6 +12,10 @@
 
 
 class ApiClientAdapter(ApiClient):
+    def __init__(self, *args, **kwargs):
+        self._token_lock = asyncio.Lock()
+        super().__init__(*args, **kwargs)
+
     async def call_api(
         self,
         method,
@@ -40,17 +46,37 @@ async def call_api(
                 post_params=post_params,
                 _request_timeout=_request_timeout,
             )
-            if response_data.status == 401:  # noqa: PLR2004 (Unauthorized status code)
-                token = await self.refresh_authorization_token()
-                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,
-                )
+            if (
+                response_data.status == 401  # noqa: PLR2004 (Unauthorized status code)
+                and url != self.configuration.host + "/token"
+            ):
+                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"]
+                    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,
+                    )
         except ApiException as e:
             raise e
 
@@ -59,7 +85,8 @@ async def call_api(
     async def refresh_authorization_token(self):
         obtain_new_token_response = await self.obtain_new_token()
         token = obtain_new_token_response.get("token")
-        self.configuration.api_key["api_key"] = token
+        self.configuration._http_config.api_key["api_key"] = token
+        self.configuration.token_update_time = time.time()
         return token
 
     async def obtain_new_token(self):

src/conductor/asyncio_client/configuration/configuration.py

@@ -4,8 +4,9 @@
 import os
 from typing import Any, Dict, Optional, Union
 
-from conductor.asyncio_client.http.configuration import \
-    Configuration as HttpConfiguration
+from conductor.asyncio_client.http.configuration import (
+    Configuration as HttpConfiguration,
+)
 
 
 class Configuration:
@@ -55,6 +56,7 @@ def __init__(
         auth_key: Optional[str] = None,
         auth_secret: Optional[str] = None,
         debug: bool = False,
+        auth_token_ttl_min: int = 45,
         # Worker properties
         default_polling_interval: Optional[float] = None,
         default_domain: Optional[str] = None,
@@ -128,10 +130,6 @@ def __init__(
         if api_key is None:
             api_key = {}
 
-        if self.auth_key and self.auth_secret:
-            # Use the auth_key as the API key for X-Authorization header
-            api_key["api_key"] = self.auth_key
-
         self.__ui_host = os.getenv("CONDUCTOR_UI_SERVER_URL")
         if self.__ui_host is None:
             self.__ui_host = self.server_url.replace("/api", "")
@@ -172,6 +170,10 @@ def __init__(
         if debug:
             self.logger.setLevel(logging.DEBUG)
 
+        # Orkes Conductor auth token properties
+        self.token_update_time = 0
+        self.auth_token_ttl_sec = auth_token_ttl_min * 60
+
     def _get_env_float(self, env_var: str, default: float) -> float:
         """Get float value from environment variable with default fallback."""
         try:
@@ -450,16 +452,13 @@ def log_level(self) -> int:
         """Get log level."""
         return self.__log_level
 
-    def apply_logging_config(self, log_format : Optional[str] = None, level = None):
+    def apply_logging_config(self, log_format: Optional[str] = None, level=None):
         """Apply logging configuration for the application."""
         if log_format is None:
             log_format = self.logger_format
         if level is None:
             level = self.__log_level
-        logging.basicConfig(
-            format=log_format,
-            level=level
-        )
+        logging.basicConfig(format=log_format, level=level)
 
     @staticmethod
     def get_logging_formatted_name(name):
@@ -474,5 +473,7 @@ def ui_host(self):
     def __getattr__(self, name: str) -> Any:
         """Delegate attribute access to underlying HTTP configuration."""
         if "_http_config" not in self.__dict__ or self._http_config is None:
-            raise AttributeError(f"'{self.__class__.__name__}' object has no attribute '{name}'")
+            raise AttributeError(
+                f"'{self.__class__.__name__}' object has no attribute '{name}'"
+            )
         return getattr(self._http_config, name)

tests/unit/api_client/test_async_api_client.py

@@ -0,0 +1,221 @@
+import asyncio
+import time
+
+import pytest
+
+from conductor.asyncio_client.adapters import ApiClient
+from conductor.asyncio_client.configuration.configuration import Configuration
+from conductor.asyncio_client.http.rest import RESTResponse
+
+
+@pytest.fixture
+def api_client():
+    configuration = Configuration(
+        server_url="http://localhost:8080/api",
+        auth_key="test_key",
+        auth_secret="test_secret",
+    )
+    return ApiClient(configuration)
+
+
+@pytest.fixture
+def mock_success_response(mocker):
+    response = mocker.Mock(spec=RESTResponse)
+    response.status = 200
+    response.data = b'{"token": "test_token"}'
+    response.read = mocker.Mock()
+    return response
+
+
+@pytest.fixture
+def mock_401_response(mocker):
+    response = mocker.Mock(spec=RESTResponse)
+    response.status = 401
+    response.data = b'{"message":"Token cannot be null or empty","error":"INVALID_TOKEN","timestamp":1758039192168}'
+    response.read = mocker.AsyncMock()
+    return response
+
+
+@pytest.mark.asyncio
+async def test_refresh_authorization_token_called_on_invalid_token(
+    mocker, api_client, mock_401_response, mock_success_response
+):
+    api_client.configuration._http_config.api_key = {}
+
+    api_client.rest_client = mocker.AsyncMock()
+    api_client.rest_client.request.side_effect = [
+        mock_401_response,
+        mock_success_response,
+    ]
+
+    mock_refresh = mocker.patch.object(
+        api_client, "refresh_authorization_token", new_callable=mocker.AsyncMock
+    )
+    mock_refresh.return_value = "new_token"
+
+    mock_obtain = mocker.patch.object(
+        api_client, "obtain_new_token", new_callable=mocker.AsyncMock
+    )
+    mock_obtain.return_value = {"token": "new_token"}
+
+    await api_client.call_api(
+        method="GET", url="http://localhost:8080/api/test", header_params={}
+    )
+
+    mock_refresh.assert_called_once()
+
+
+@pytest.mark.asyncio
+async def test_refresh_authorization_token_called_on_expired_token(
+    mocker, api_client, mock_401_response, mock_success_response
+):
+    current_time = time.time()
+    api_client.configuration.token_update_time = current_time - 3600
+    api_client.configuration.auth_token_ttl_sec = 1800
+    api_client.configuration._http_config.api_key = {"api_key": "old_token"}
+
+    api_client.rest_client = mocker.AsyncMock()
+    api_client.rest_client.request.side_effect = [
+        mock_401_response,
+        mock_success_response,
+    ]
+
+    mock_refresh = mocker.patch.object(
+        api_client, "refresh_authorization_token", new_callable=mocker.AsyncMock
+    )
+    mock_refresh.return_value = "new_token"
+
+    mock_obtain = mocker.patch.object(
+        api_client, "obtain_new_token", new_callable=mocker.AsyncMock
+    )
+    mock_obtain.return_value = {"token": "new_token"}
+
+    await api_client.call_api(
+        method="GET", url="http://localhost:8080/api/test", header_params={}
+    )
+
+    mock_refresh.assert_called_once()
+
+
+@pytest.mark.asyncio
+async def test_token_lock_prevents_concurrent_refresh(
+    mocker, api_client, mock_401_response, mock_success_response
+):
+    api_client.configuration._http_config.api_key = {}
+
+    refresh_calls = []
+
+    async def mock_refresh():
+        refresh_calls.append(time.time())
+        await asyncio.sleep(0.1)
+        return "new_token"
+
+    mocker.patch.object(
+        api_client, "refresh_authorization_token", side_effect=mock_refresh
+    )
+
+    mock_obtain = mocker.patch.object(
+        api_client, "obtain_new_token", new_callable=mocker.AsyncMock
+    )
+    mock_obtain.return_value = {"token": "new_token"}
+
+    api_client.rest_client = mocker.AsyncMock()
+    api_client.rest_client.request.side_effect = [
+        mock_401_response,
+        mock_success_response,
+        mock_401_response,
+        mock_success_response,
+    ]
+
+    tasks = [
+        api_client.call_api(
+            method="GET",
+            url="http://localhost:8080/api/test1",
+            header_params={},
+        ),
+        api_client.call_api(
+            method="GET",
+            url="http://localhost:8080/api/test2",
+            header_params={},
+        ),
+    ]
+
+    await asyncio.gather(*tasks)
+
+    assert len(refresh_calls) == 1
+
+
+@pytest.mark.asyncio
+async def test_no_refresh_when_token_valid_and_not_expired(
+    mocker, api_client, mock_success_response
+):
+    current_time = time.time()
+    api_client.configuration.token_update_time = current_time - 100
+    api_client.configuration.auth_token_ttl_sec = 1800
+    api_client.configuration._http_config.api_key = {"api_key": "valid_token"}
+
+    api_client.rest_client = mocker.AsyncMock()
+    api_client.rest_client.request.return_value = mock_success_response
+
+    mock_refresh = mocker.patch.object(
+        api_client, "refresh_authorization_token", new_callable=mocker.AsyncMock
+    )
+
+    await api_client.call_api(
+        method="GET", url="http://localhost:8080/api/test", header_params={}
+    )
+
+    mock_refresh.assert_not_called()
+
+
+@pytest.mark.asyncio
+async def test_no_refresh_for_token_endpoint(mocker, api_client, mock_401_response):
+    api_client.configuration._http_config.api_key = {}
+
+    api_client.rest_client = mocker.AsyncMock()
+    api_client.rest_client.request.return_value = mock_401_response
+
+    mock_refresh = mocker.patch.object(
+        api_client, "refresh_authorization_token", new_callable=mocker.AsyncMock
+    )
+
+    await api_client.call_api(
+        method="POST", url="http://localhost:8080/api/token", header_params={}
+    )
+
+    mock_refresh.assert_not_called()
+
+
+@pytest.mark.asyncio
+async def test_401_response_triggers_retry_with_new_token(
+    mocker, api_client, mock_401_response, mock_success_response
+):
+    api_client.configuration._http_config.api_key = {}
+
+    api_client.rest_client = mocker.AsyncMock()
+    api_client.rest_client.request.side_effect = [
+        mock_401_response,
+        mock_success_response,
+    ]
+
+    mock_refresh = mocker.patch.object(
+        api_client, "refresh_authorization_token", new_callable=mocker.AsyncMock
+    )
+    mock_refresh.return_value = "new_token"
+
+    mock_obtain = mocker.patch.object(
+        api_client, "obtain_new_token", new_callable=mocker.AsyncMock
+    )
+    mock_obtain.return_value = {"token": "new_token"}
+
+    header_params = {}
+    await api_client.call_api(
+        method="GET",
+        url="http://localhost:8080/api/test",
+        header_params=header_params,
+    )
+
+    assert api_client.rest_client.request.call_count == 2
+
+    second_call_args = api_client.rest_client.request.call_args_list[1]
+    assert second_call_args[1]["headers"]["X-Authorization"] == "new_token"

tests/unit/configuration/test_async_configuration.py

@@ -0,0 +1,25 @@
+import pytest
+
+from conductor.asyncio_client.configuration import Configuration
+
+
+def test_initialization_default(monkeypatch):
+    monkeypatch.setenv("CONDUCTOR_SERVER_URL", "http://localhost:8080/api")
+    configuration = Configuration()
+    assert configuration.host == "http://localhost:8080/api"
+
+
+def test_initialization_with_base_url():
+    configuration = Configuration(server_url="https://play.orkes.io/api")
+    assert configuration.host == "https://play.orkes.io/api"
+
+
+def test_missed_http_config():
+    configuration = Configuration()
+    configuration._http_config = None
+    with pytest.raises(AttributeError) as ctx:
+        _ = configuration.api_key
+        assert (
+            f"'{Configuration.__class__.__name__}' object has no attribute 'api_key'"
+            in ctx.value
+        )