From 28759f7d7028cc185d54f2232388358da579cca5 Mon Sep 17 00:00:00 2001 From: Muhammad Zohaib Date: Sat, 31 Oct 2026 23:31:24 +0300 Subject: [PATCH 1/2] feat(temporal): add workload identity authentication --- agentex/pyproject.toml | 1 + .../src/adapters/temporal/client_factory.py | 3 +- agentex/src/utils/temporal_auth.py | 112 +++++++++++ agentex/src/utils/temporal_client.py | 16 +- .../tests/unit/utils/test_temporal_auth.py | 188 ++++++++++++++++++ uv.lock | 57 ++++++ 6 files changed, 370 insertions(+), 7 deletions(-) create mode 100644 agentex/src/utils/temporal_auth.py create mode 100644 agentex/tests/unit/utils/test_temporal_auth.py diff --git a/agentex/pyproject.toml b/agentex/pyproject.toml index 5e5ba384..fe1d0d1a 100644 --- a/agentex/pyproject.toml +++ b/agentex/pyproject.toml @@ -9,6 +9,7 @@ dependencies = [ "fastapi>=0.115.0,<0.137.0", # <0.137.0 keeps app.routes flat; 0.137.0's _IncludedRouter (no .path) crashes injected OTel instrumentation on OPTIONS preflights -> 500 "litellm>=1.84.0,<2", # >=1.84.0 clears CVE-2026-49468 (CRITICAL auth bypass) "python-dotenv>=1.2.2,<2", + "azure-identity>=1.16.1,<2.0.0", "temporalio>=1.18.0,<2", "uvicorn[standard]>=0.35.0,<0.36", "redis>=5.1.0,<6", diff --git a/agentex/src/adapters/temporal/client_factory.py b/agentex/src/adapters/temporal/client_factory.py index fd48551a..1a5981d2 100644 --- a/agentex/src/adapters/temporal/client_factory.py +++ b/agentex/src/adapters/temporal/client_factory.py @@ -22,6 +22,7 @@ from src.config.environment_variables import EnvironmentVariables from src.utils.logging import make_logger from src.utils.observability import temporal_client_interceptors +from src.utils.temporal_auth import connect_temporal_client logger = make_logger(__name__) @@ -137,7 +138,7 @@ async def create_client( connect_options["runtime"] = runtime # Create the client - client = await Client.connect(**connect_options) + client = await connect_temporal_client(connect_options) logger.info( f"Successfully created Temporal client for address: {temporal_address}" ) diff --git a/agentex/src/utils/temporal_auth.py b/agentex/src/utils/temporal_auth.py new file mode 100644 index 00000000..e556ede4 --- /dev/null +++ b/agentex/src/utils/temporal_auth.py @@ -0,0 +1,112 @@ +"""Renewable Azure Workload Identity authentication for Temporal clients.""" + +from __future__ import annotations + +import asyncio +import logging +import os +import time +from collections.abc import Awaitable, Callable +from dataclasses import dataclass, field +from typing import Any, Protocol + +from azure.core.credentials import AccessToken +from azure.identity.aio import WorkloadIdentityCredential +from temporalio.client import Client + +logger = logging.getLogger(__name__) + +_AUTH_SCOPE_ENV = "TEMPORAL_AUTH_SCOPE" +_REFRESH_MARGIN_SECONDS = 300 +_MIN_REFRESH_DELAY_SECONDS = 30 +_RETRY_DELAY_SECONDS = 30 + + +class AsyncTokenCredential(Protocol): + """Credential surface needed by the refresher.""" + + async def get_token(self, *scopes: str) -> AccessToken: ... + + async def close(self) -> None: ... + + +@dataclass(eq=False) +class TemporalTokenSession: + """Keep a Temporal client's bearer token current for the process lifetime.""" + + client: Client + credential: AsyncTokenCredential + scope: str + token: AccessToken + now: Callable[[], float] = time.time + sleep: Callable[[float], Awaitable[None]] = asyncio.sleep + _task: asyncio.Task[None] | None = field(default=None, init=False) + + def start(self) -> None: + self._task = asyncio.create_task(self._run()) + _ACTIVE_SESSIONS.add(self) + self._task.add_done_callback(lambda _: _ACTIVE_SESSIONS.discard(self)) + + async def refresh_once(self) -> bool: + """Replace the token only after a valid refresh succeeds.""" + try: + refreshed = await self.credential.get_token(self.scope) + if refreshed.expires_on <= self.now(): + raise ValueError("Azure returned an expired Temporal access token") + except Exception: + log = ( + logger.error if self.token.expires_on <= self.now() else logger.warning + ) + log( + "Temporal access token refresh failed; retaining the current token", + exc_info=True, + ) + return False + + self.client.api_key = refreshed.token + self.token = refreshed + return True + + async def _run(self) -> None: + try: + while True: + refresh_in = max( + _MIN_REFRESH_DELAY_SECONDS, + self.token.expires_on - self.now() - _REFRESH_MARGIN_SECONDS, + ) + await self.sleep(refresh_in) + if not await self.refresh_once(): + await self.sleep(_RETRY_DELAY_SECONDS) + finally: + await self.credential.close() + + +_ACTIVE_SESSIONS: set[TemporalTokenSession] = set() + + +async def connect_temporal_client(connect_options: dict[str, Any]) -> Client: + """Connect anonymously by default, or with renewable Workload Identity.""" + scope = os.environ.get(_AUTH_SCOPE_ENV) + if not scope: + return await Client.connect(**connect_options) + + credential = WorkloadIdentityCredential() + try: + token = await credential.get_token(scope) + if token.expires_on <= time.time(): + raise ValueError("Azure returned an expired Temporal access token") + authenticated_options = dict(connect_options) + if authenticated_options.get("tls") is None: + authenticated_options["tls"] = False + client = await Client.connect(api_key=token.token, **authenticated_options) + except Exception: + await credential.close() + raise + + TemporalTokenSession( + client=client, + credential=credential, + scope=scope, + token=token, + ).start() + return client diff --git a/agentex/src/utils/temporal_client.py b/agentex/src/utils/temporal_client.py index bdc8bdbb..e73b3a3d 100644 --- a/agentex/src/utils/temporal_client.py +++ b/agentex/src/utils/temporal_client.py @@ -14,6 +14,8 @@ ) from temporalio.runtime import OpenTelemetryConfig, Runtime, TelemetryConfig +from src.utils.temporal_auth import connect_temporal_client + class DateTimeJSONEncoder(AdvancedJSONEncoder): def default(self, o: Any) -> Any: @@ -53,16 +55,18 @@ def __init__(self) -> None: async def get_temporal_client(temporal_address: str, metrics_url: str = None) -> Client: if not metrics_url: - client = await Client.connect( - target_host=temporal_address, data_converter=custom_data_converter + client = await connect_temporal_client( + {"target_host": temporal_address, "data_converter": custom_data_converter} ) else: runtime = Runtime( telemetry=TelemetryConfig(metrics=OpenTelemetryConfig(url=metrics_url)) ) - client = await Client.connect( - target_host=temporal_address, - data_converter=custom_data_converter, - runtime=runtime, + client = await connect_temporal_client( + { + "target_host": temporal_address, + "data_converter": custom_data_converter, + "runtime": runtime, + } ) return client diff --git a/agentex/tests/unit/utils/test_temporal_auth.py b/agentex/tests/unit/utils/test_temporal_auth.py new file mode 100644 index 00000000..43bc87e3 --- /dev/null +++ b/agentex/tests/unit/utils/test_temporal_auth.py @@ -0,0 +1,188 @@ +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest +from azure.core.credentials import AccessToken +from src.utils.temporal_auth import TemporalTokenSession, connect_temporal_client +from temporalio.client import TLSConfig + +_SCOPE = "api:" + "//example.invalid/.default" +_INITIAL = "fixture" + "-initial" +_RENEWED = "fixture" + "-renewed" + + +class FakeCredential: + def __init__(self, result): + self.result = result + self.closed = False + + async def get_token(self, _scope): + if isinstance(self.result, Exception): + raise self.result + return self.result + + async def close(self): + self.closed = True + + +async def test_connect_without_scope_preserves_legacy_path(monkeypatch): + monkeypatch.delenv("TEMPORAL_AUTH_SCOPE", raising=False) + client = MagicMock() + + with ( + patch("src.utils.temporal_auth.WorkloadIdentityCredential") as credential, + patch( + "src.utils.temporal_auth.Client.connect", new=AsyncMock(return_value=client) + ) as connect, + ): + result = await connect_temporal_client({"target_host": "temporal:7233"}) + + assert result is client + credential.assert_not_called() + connect.assert_awaited_once_with(target_host="temporal:7233") + + +async def test_connect_initial_token_failure_blocks_startup(monkeypatch): + monkeypatch.setenv("TEMPORAL_AUTH_SCOPE", _SCOPE) + credential = FakeCredential(RuntimeError("identity unavailable")) + + with ( + patch( + "src.utils.temporal_auth.WorkloadIdentityCredential", + return_value=credential, + ), + patch("src.utils.temporal_auth.Client.connect", new=AsyncMock()) as connect, + pytest.raises(RuntimeError, match="identity unavailable"), + ): + await connect_temporal_client({"target_host": "temporal:7233"}) + + connect.assert_not_awaited() + assert credential.closed is True + + +async def test_connect_expired_initial_token_blocks_startup(monkeypatch): + monkeypatch.setenv("TEMPORAL_AUTH_SCOPE", _SCOPE) + credential = FakeCredential(AccessToken(_INITIAL, 999)) + + with ( + patch("src.utils.temporal_auth.time.time", return_value=1_000), + patch( + "src.utils.temporal_auth.WorkloadIdentityCredential", + return_value=credential, + ), + patch("src.utils.temporal_auth.Client.connect", new=AsyncMock()) as connect, + pytest.raises(ValueError, match="expired Temporal access token"), + ): + await connect_temporal_client({"target_host": "temporal:7233"}) + + connect.assert_not_awaited() + assert credential.closed is True + + +async def test_connect_sends_initial_token_over_plaintext(monkeypatch): + monkeypatch.setenv("TEMPORAL_AUTH_SCOPE", _SCOPE) + access = AccessToken(_INITIAL, 2_000) + credential = FakeCredential(access) + client = MagicMock() + + with ( + patch("src.utils.temporal_auth.time.time", return_value=1_000), + patch( + "src.utils.temporal_auth.WorkloadIdentityCredential", + return_value=credential, + ), + patch( + "src.utils.temporal_auth.Client.connect", new=AsyncMock(return_value=client) + ) as connect, + patch.object(TemporalTokenSession, "start") as start, + ): + result = await connect_temporal_client({"target_host": "temporal:7233"}) + + assert result is client + connect.assert_awaited_once_with( + target_host="temporal:7233", tls=False, api_key=_INITIAL + ) + start.assert_called_once() + + +async def test_connect_preserves_explicit_tls(monkeypatch): + monkeypatch.setenv("TEMPORAL_AUTH_SCOPE", _SCOPE) + access = AccessToken(_INITIAL, 2_000) + credential = FakeCredential(access) + client = MagicMock() + tls = TLSConfig(server_root_ca_cert=b"fixture-ca") + + with ( + patch("src.utils.temporal_auth.time.time", return_value=1_000), + patch( + "src.utils.temporal_auth.WorkloadIdentityCredential", + return_value=credential, + ), + patch( + "src.utils.temporal_auth.Client.connect", new=AsyncMock(return_value=client) + ) as connect, + patch.object(TemporalTokenSession, "start") as start, + ): + result = await connect_temporal_client( + {"target_host": "temporal:7233", "tls": tls} + ) + + assert result is client + connect.assert_awaited_once_with( + target_host="temporal:7233", tls=tls, api_key=_INITIAL + ) + start.assert_called_once() + + +async def test_refresh_replaces_the_client_token(): + client = MagicMock(api_key=_INITIAL) + renewed = AccessToken(_RENEWED, 3_000) + credential = FakeCredential(renewed) + initial = AccessToken(_INITIAL, 2_000) + session = TemporalTokenSession( + client=client, + credential=credential, + scope=_SCOPE, + token=initial, + now=lambda: 1_000, + ) + + refreshed = await session.refresh_once() + + assert refreshed is True + assert client.api_key == _RENEWED + + +async def test_refresh_failure_never_clears_authentication(): + client = MagicMock(api_key=_INITIAL) + credential = FakeCredential(RuntimeError("identity unavailable")) + initial = AccessToken(_INITIAL, 900) + session = TemporalTokenSession( + client=client, + credential=credential, + scope=_SCOPE, + token=initial, + now=lambda: 1_000, + ) + + refreshed = await session.refresh_once() + + assert refreshed is False + assert client.api_key == _INITIAL + + +async def test_refresh_expired_replacement_never_clears_authentication(): + client = MagicMock(api_key=_INITIAL) + credential = FakeCredential(AccessToken(_RENEWED, 999)) + initial = AccessToken(_INITIAL, 2_000) + session = TemporalTokenSession( + client=client, + credential=credential, + scope=_SCOPE, + token=initial, + now=lambda: 1_000, + ) + + refreshed = await session.refresh_once() + + assert refreshed is False + assert client.api_key == _INITIAL diff --git a/uv.lock b/uv.lock index d62eb5e0..1d3a19d2 100644 --- a/uv.lock +++ b/uv.lock @@ -63,6 +63,7 @@ dependencies = [ { name = "aiohttp" }, { name = "alembic" }, { name = "asyncpg" }, + { name = "azure-identity" }, { name = "cryptography" }, { name = "datadog" }, { name = "ddtrace" }, @@ -123,6 +124,7 @@ requires-dist = [ { name = "aiohttp", specifier = ">=3.10.9,<4" }, { name = "alembic", specifier = ">=1.13.3,<2" }, { name = "asyncpg", specifier = ">=0.29.0,<0.30" }, + { name = "azure-identity", specifier = ">=1.16.1,<2.0.0" }, { name = "cryptography", specifier = ">=50.0.0" }, { name = "datadog", specifier = ">=0.52.1" }, { name = "ddtrace", specifier = ">=4.8.2" }, @@ -402,6 +404,35 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/77/06/bb80f5f86020c4551da315d78b3ab75e8228f89f0162f2c3a819e407941a/attrs-25.3.0-py3-none-any.whl", hash = "sha256:427318ce031701fea540783410126f03899a97ffc6f61596ad581ac2e40e3bc3", size = 63815, upload-time = "2025-03-13T11:10:21.14Z" }, ] +[[package]] +name = "azure-core" +version = "1.41.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "requests" }, + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/a6/f3/b416179e408990df5db0d516283022dde0f5d0111d98c1a848e41853e81c/azure_core-1.41.0.tar.gz", hash = "sha256:f46ff5dfcd230f25cf1c19e8a34b8dc08a337b2503e268bb600a16c00db8ad5a", size = 381042, upload-time = "2026-05-07T23:30:54.302Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/5b/db/325c6d7312d2200251c52323878281045aaffcb5586612296484e4280eaa/azure_core-1.41.0-py3-none-any.whl", hash = "sha256:522b4011e8180b1a3dcd2024396a4e7fe9ac37fb8597db47163d230b5efe892d", size = 220920, upload-time = "2026-05-07T23:30:56.357Z" }, +] + +[[package]] +name = "azure-identity" +version = "1.27.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "azure-core" }, + { name = "cryptography" }, + { name = "msal" }, + { name = "msal-extensions" }, + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/cf/33/0b881d0093d64c965bb20fcf71b806c7097881b544a2a55d519cdcf3bab9/azure_identity-1.27.0.tar.gz", hash = "sha256:6a4c50e25e8e9a47b324a95ad024897e56cb3c6241e0094e93b126a2bf71edc5", size = 331360, upload-time = "2026-10-26T21:17:39.095Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/29/1e/6a2fe486576f41884a120ef7b2c3774f334f9143402fb73650c67db94ddd/azure_identity-1.27.0-py3-none-any.whl", hash = "sha256:92fd2a93c0d2f777dac6c04e5fd7d5734f6643582b983544ef66b4575844e46d", size = 203474, upload-time = "2026-10-26T21:17:40.92Z" }, +] + [[package]] name = "bytecode" version = "0.17.0" @@ -1522,6 +1553,32 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/b3/38/89ba8ad64ae25be8de66a6d463314cf1eb366222074cfda9ee839c56a4b4/mdurl-0.1.2-py3-none-any.whl", hash = "sha256:84008a41e51615a49fc9966191ff91509e3c40b939176e643fd50a5c2196b8f8", size = 9979, upload-time = "2022-08-14T12:40:09.779Z" }, ] +[[package]] +name = "msal" +version = "1.39.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "cryptography" }, + { name = "pyjwt", extra = ["crypto"] }, + { name = "requests" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/bb/85/747d28986a44b715cc51cf57aa021c8fd94e9614064261ceec77320e8b0b/msal-1.39.0.tar.gz", hash = "sha256:6ab7de335e6d7f5717e2c7e1dbf86e4dda2f6acf3c56773b78dc53ebc6395b5f", size = 205360, upload-time = "2026-09-17T16:07:45.1Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/a4/32/178457385925ce301c7b3232bd85665d4df57da39568b62dcfc74031873f/msal-1.39.0-py3-none-any.whl", hash = "sha256:2d2577886906cd7293850dffa2da29119966c213bfc6ec0cecf8bf7621e1ca77", size = 131629, upload-time = "2026-09-17T16:07:46.722Z" }, +] + +[[package]] +name = "msal-extensions" +version = "1.3.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "msal" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/01/99/5d239b6156eddf761a636bded1118414d161bd6b7b37a9335549ed159396/msal_extensions-1.3.1.tar.gz", hash = "sha256:c5b0fd10f65ef62b5f1d62f4251d51cbcaf003fcedae8c91b040a488614be1a4", size = 23315, upload-time = "2025-03-14T23:51:03.902Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/5e/75/bd9b7bb966668920f06b200e84454c8f3566b102183bc55c5473d96cb2b9/msal_extensions-1.3.1-py3-none-any.whl", hash = "sha256:96d3de4d034504e969ac5e85bae8106c8373b5c6568e4c8fa7af2eca9dbe6bca", size = 20583, upload-time = "2025-03-14T23:51:03.016Z" }, +] + [[package]] name = "multidict" version = "6.6.4" From faa3420ee645e8b1d5bf7b7af4715b118e66e76c Mon Sep 17 00:00:00 2001 From: Muhammad Zohaib Date: Sun, 1 Nov 2026 01:37:09 +0300 Subject: [PATCH 2/2] fix(temporal): correct token refresh retries --- agentex/src/utils/temporal_auth.py | 18 ++++--- .../tests/unit/utils/test_temporal_auth.py | 47 +++++++++++++++++++ 2 files changed, 59 insertions(+), 6 deletions(-) diff --git a/agentex/src/utils/temporal_auth.py b/agentex/src/utils/temporal_auth.py index e556ede4..252ceee6 100644 --- a/agentex/src/utils/temporal_auth.py +++ b/agentex/src/utils/temporal_auth.py @@ -69,14 +69,20 @@ async def refresh_once(self) -> bool: async def _run(self) -> None: try: + delay = max( + _MIN_REFRESH_DELAY_SECONDS, + self.token.expires_on - self.now() - _REFRESH_MARGIN_SECONDS, + ) while True: - refresh_in = max( - _MIN_REFRESH_DELAY_SECONDS, - self.token.expires_on - self.now() - _REFRESH_MARGIN_SECONDS, + await self.sleep(delay) + delay = ( + max( + _MIN_REFRESH_DELAY_SECONDS, + self.token.expires_on - self.now() - _REFRESH_MARGIN_SECONDS, + ) + if await self.refresh_once() + else _RETRY_DELAY_SECONDS ) - await self.sleep(refresh_in) - if not await self.refresh_once(): - await self.sleep(_RETRY_DELAY_SECONDS) finally: await self.credential.close() diff --git a/agentex/tests/unit/utils/test_temporal_auth.py b/agentex/tests/unit/utils/test_temporal_auth.py index 43bc87e3..a8cf31d4 100644 --- a/agentex/tests/unit/utils/test_temporal_auth.py +++ b/agentex/tests/unit/utils/test_temporal_auth.py @@ -24,6 +24,24 @@ async def close(self): self.closed = True +class RetryCredential: + def __init__(self, results, now): + self.results = iter(results) + self.now = now + self.call_times = [] + self.closed = False + + async def get_token(self, _scope): + self.call_times.append(self.now()) + result = next(self.results) + if isinstance(result, Exception): + raise result + return result + + async def close(self): + self.closed = True + + async def test_connect_without_scope_preserves_legacy_path(monkeypatch): monkeypatch.delenv("TEMPORAL_AUTH_SCOPE", raising=False) client = MagicMock() @@ -186,3 +204,32 @@ async def test_refresh_expired_replacement_never_clears_authentication(): assert refreshed is False assert client.api_key == _INITIAL + + +async def test_refresh_failure_retries_after_one_retry_delay(): + clock = {"now": 0} + initial = AccessToken(_INITIAL, 345) + credential = RetryCredential( + [RuntimeError("identity unavailable"), AccessToken(_RENEWED, 1_000)], + lambda: clock["now"], + ) + + async def advance(delay): + clock["now"] += delay + if len(credential.call_times) == 2: + raise StopAsyncIteration + + session = TemporalTokenSession( + client=MagicMock(api_key=_INITIAL), + credential=credential, + scope=_SCOPE, + token=initial, + now=lambda: clock["now"], + sleep=advance, + ) + + with pytest.raises(StopAsyncIteration): + await session._run() + + assert credential.call_times == [45, 75] + assert credential.closed is True