feat(temporal): add workload identity authentication

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

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

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)

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: