feat(temporal): add workload identity authentication

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",

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}"
             )

agentex/src/utils/temporal_auth.py

@@ -0,0 +1,118 @@
+"""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:
+            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

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