diff --git a/adk/pyproject.toml b/adk/pyproject.toml index c81346591..119b79b71 100644 --- a/adk/pyproject.toml +++ b/adk/pyproject.toml @@ -34,6 +34,7 @@ dependencies = [ "uvicorn>=0.31.1", "aiohttp>=3.10.10,<4", # Temporal workflows + "azure-identity>=1.16.1,<2.0.0", "temporalio>=1.26.0,<1.35", "cloudpickle>=3.1.1", # Async streaming infra diff --git a/src/agentex/lib/core/clients/temporal/auth.py b/src/agentex/lib/core/clients/temporal/auth.py new file mode 100644 index 000000000..509e43838 --- /dev/null +++ b/src/agentex/lib/core/clients/temporal/auth.py @@ -0,0 +1,111 @@ +"""Renewable Azure Workload Identity authentication for Temporal clients.""" + +from __future__ import annotations + +import os +import time +import asyncio +import logging +from typing import Any, Protocol +from dataclasses import field, dataclass +from collections.abc import Callable, Awaitable + +from temporalio.client import Client +from azure.identity.aio import WorkloadIdentityCredential +from azure.core.credentials import AccessToken + +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): + 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: + delay = max( + _MIN_REFRESH_DELAY_SECONDS, + self.token.expires_on - self.now() - _REFRESH_MARGIN_SECONDS, + ) + while True: + 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 + ) + 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/src/agentex/lib/core/clients/temporal/utils.py b/src/agentex/lib/core/clients/temporal/utils.py index 15b08cec6..0e41fb683 100644 --- a/src/agentex/lib/core/clients/temporal/utils.py +++ b/src/agentex/lib/core/clients/temporal/utils.py @@ -10,6 +10,7 @@ from temporalio.contrib.pydantic import pydantic_data_converter from agentex.lib.core.tracing.temporal import temporal_tracing_interceptors +from agentex.lib.core.clients.temporal.auth import connect_temporal_client # class DateTimeJSONEncoder(AdvancedJSONEncoder): # def default(self, o: Any) -> Any: @@ -151,10 +152,7 @@ async def get_temporal_client( dc = dataclasses.replace(dc, payload_codec=payload_codec) connect_kwargs["data_converter"] = dc - if not metrics_url: - client = await Client.connect(**connect_kwargs) - else: + if metrics_url: runtime = Runtime(telemetry=TelemetryConfig(metrics=OpenTelemetryConfig(url=metrics_url))) connect_kwargs["runtime"] = runtime - client = await Client.connect(**connect_kwargs) - return client + return await connect_temporal_client(connect_kwargs) diff --git a/src/agentex/lib/core/temporal/workers/worker.py b/src/agentex/lib/core/temporal/workers/worker.py index ba8f87de5..8256cba03 100644 --- a/src/agentex/lib/core/temporal/workers/worker.py +++ b/src/agentex/lib/core/temporal/workers/worker.py @@ -33,6 +33,7 @@ from agentex.lib.environment_variables import EnvironmentVariables from agentex.lib.core.tracing.span_queue import shutdown_default_span_queue from agentex.lib.core.compat.version_guard import assert_backend_compatible +from agentex.lib.core.clients.temporal.auth import connect_temporal_client from agentex.lib.core.observability.sgp_obs_setup import init_sgp_obs, shutdown_sgp_obs from agentex.lib.core.tracing.tracing_processor_manager import shutdown_sync_tracing_processors @@ -147,22 +148,23 @@ async def get_temporal_client( dc = dataclasses.replace(dc, payload_codec=payload_codec) connect_kwargs["data_converter"] = dc - if not metrics_url: - client = await Client.connect(**connect_kwargs) - else: - runtime = Runtime(telemetry=TelemetryConfig(metrics=OpenTelemetryConfig( - url=metrics_url, - headers=metrics_headers or {}, - http=metrics_use_http, - metric_temporality=( - OpenTelemetryMetricTemporality.DELTA - if metrics_temporality_delta - else OpenTelemetryMetricTemporality.CUMULATIVE - ), - ))) + if metrics_url: + runtime = Runtime( + telemetry=TelemetryConfig( + metrics=OpenTelemetryConfig( + url=metrics_url, + headers=metrics_headers or {}, + http=metrics_use_http, + metric_temporality=( + OpenTelemetryMetricTemporality.DELTA + if metrics_temporality_delta + else OpenTelemetryMetricTemporality.CUMULATIVE + ), + ) + ) + ) connect_kwargs["runtime"] = runtime - client = await Client.connect(**connect_kwargs) - return client + return await connect_temporal_client(connect_kwargs) class AgentexWorker: diff --git a/tests/lib/test_temporal_auth.py b/tests/lib/test_temporal_auth.py new file mode 100644 index 000000000..8be0ddfcd --- /dev/null +++ b/tests/lib/test_temporal_auth.py @@ -0,0 +1,226 @@ +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest +from temporalio.service import TLSConfig +from azure.core.credentials import AccessToken + +from agentex.lib.core.clients.temporal.auth import ( + TemporalTokenSession, + connect_temporal_client, +) + +_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 + + +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() + + with patch("agentex.lib.core.clients.temporal.auth.WorkloadIdentityCredential") as credential: + with patch( + "agentex.lib.core.clients.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( + "agentex.lib.core.clients.temporal.auth.WorkloadIdentityCredential", + return_value=credential, + ): + with patch("agentex.lib.core.clients.temporal.auth.Client.connect", new=AsyncMock()) as connect: + with 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( + "agentex.lib.core.clients.temporal.auth.WorkloadIdentityCredential", + return_value=credential, + ): + with patch("agentex.lib.core.clients.temporal.auth.Client.connect", new=AsyncMock()) as connect: + with patch("agentex.lib.core.clients.temporal.auth.time.time", return_value=1_000): + with 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( + "agentex.lib.core.clients.temporal.auth.WorkloadIdentityCredential", + return_value=credential, + ): + with patch( + "agentex.lib.core.clients.temporal.auth.Client.connect", + new=AsyncMock(return_value=client), + ) as connect: + with patch("agentex.lib.core.clients.temporal.auth.time.time", return_value=1_000): + with 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( + "agentex.lib.core.clients.temporal.auth.WorkloadIdentityCredential", + return_value=credential, + ): + with patch( + "agentex.lib.core.clients.temporal.auth.Client.connect", + new=AsyncMock(return_value=client), + ) as connect: + with patch("agentex.lib.core.clients.temporal.auth.time.time", return_value=1_000): + with 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 + + +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 diff --git a/uv.lock b/uv.lock index f8e8bf9a3..910b6eaa6 100644 --- a/uv.lock +++ b/uv.lock @@ -96,6 +96,7 @@ source = { editable = "adk" } dependencies = [ { name = "agentex-client" }, { name = "aiohttp" }, + { name = "azure-identity" }, { name = "claude-agent-sdk" }, { name = "cloudpickle" }, { name = "ddtrace" }, @@ -132,6 +133,7 @@ dependencies = [ requires-dist = [ { name = "agentex-client", editable = "." }, { name = "aiohttp", specifier = ">=3.10.10,<4" }, + { name = "azure-identity", specifier = ">=1.16.1,<2.0.0" }, { name = "claude-agent-sdk", specifier = ">=0.1.0" }, { name = "cloudpickle", specifier = ">=3.1.1" }, { name = "ddtrace", specifier = ">=3.13.0" }, @@ -320,6 +322,35 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/3a/2a/7cc015f5b9f5db42b7d48157e23356022889fc354a2813c15934b7cb5c0e/attrs-25.4.0-py3-none-any.whl", hash = "sha256:adcf7e2a1fb3b36ac48d97835bb6d8ade15b8dcce26aba8bf1d14847b57a3373", size = 67615, upload-time = "2025-10-06T13:54:43.17Z" }, ] +[[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" @@ -1537,6 +1568,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.7.0"