Add configurable Azure Managed token resource IDs

docs/getting-started.md

@@ -1,5 +1,92 @@
 # Getting Started
 
+## Configure Azure Managed Authentication
+
+Pass an Azure `TokenCredential` to the Azure Managed client or worker. Use
+`resource_id` to override the resource audience for token requests:
+
+```python
+from azure.identity import AzureAuthorityHosts, DefaultAzureCredential
+from durabletask.azuremanaged import DurableTaskSchedulerClient
+
+credential = DefaultAzureCredential(
+    authority=AzureAuthorityHosts.AZURE_GOVERNMENT,
+)
+
+client = DurableTaskSchedulerClient(
+    host_address="https://myaccount.usgovvirginia.durabletask.azure.us",
+    taskhub="my-task-hub",
+    token_credential=credential,
+    resource_id="https://durabletask.azure.us",
+)
+```
+
+The same `resource_id` parameter is available on `DurableTaskSchedulerWorker`,
+`AsyncDurableTaskSchedulerClient`, and the preview `SandboxActivitiesClient` and
+`SandboxWorker`. The async client requires an async credential, such as
+`azure.identity.aio.DefaultAzureCredential`. Sandbox workers continue to use
+their runtime-injected endpoint and managed identity.
+
+An explicit resource ID takes precedence over `REGION_NAME`. If `resource_id`
+is `None` or `""`, the SDK resolves the default when the client or worker is
+constructed:
+
+| `REGION_NAME` | Default resource ID |
+| --- | --- |
+| Starts with `usgov` or `usdod` (case-insensitive) | `https://durabletask.azure.us` |
+| Any other value, empty, or unset | `https://durabletask.io` |
+
+For example, to select the government default without passing `resource_id`:
+
+Bash:
+
+```bash
+export REGION_NAME=usgovvirginia
+```
+
+PowerShell:
+
+```powershell
+$env:REGION_NAME = "usgovvirginia"
+```
+
+The SDK trims surrounding whitespace and trailing slashes, removes an existing
+`/.default` suffix (case-insensitive), and appends `/.default` for the token
+request. For example, both `https://durabletask.azure.us/` and
+`https://durabletask.azure.us/.default` request
+`https://durabletask.azure.us/.default`. Nonempty inputs that become empty after
+normalization, such as whitespace, `///`, or `/.default`, raise `ValueError`.
+
+### Credential Ownership and Authority Host
+
+`DurableTaskSchedulerClient`, `AsyncDurableTaskSchedulerClient`,
+`SandboxActivitiesClient`, and `DurableTaskSchedulerWorker` require a
+`token_credential` argument. They use the caller-created credential directly;
+they do not construct `DefaultAzureCredential` or another credential on the
+caller's behalf. Passing `None` disables SDK token authentication rather than
+creating a default credential. A caller-supplied channel remains responsible
+for its own authentication.
+
+Credentials may come from Azure Identity or implement the Azure Core
+`TokenCredential` / `AsyncTokenCredential` protocol themselves. Token requests
+use `get_token(scope)`; the protocol has no supported per-request authority
+override. Configure `authority` when creating a compatible Azure Identity
+credential, as in the example above. Omitting it preserves the credential's
+default behavior, including `AZURE_AUTHORITY_HOST` where applicable.
+
+The preview `SandboxWorker` is the only Azure Managed runtime path that creates
+an Azure Identity credential internally. It creates `ManagedIdentityCredential`
+using the injected `DTS_UMI_CLIENT_ID`. Managed identities use their hosting
+environment's identity endpoint and ignore the authority setting, so this path
+does not require an authority-host option either.
+
+> [!NOTE]
+> The resource audience, service endpoint, and credential authority/cloud are
+> separate settings. Neither `resource_id` nor `REGION_NAME` changes the endpoint
+> or the credential's authority. Configure the credential and any underlying
+> developer tools for the target cloud separately. Other clouds or custom
+> audiences require an explicit resource ID; they are not inferred from the endpoint.
+
 ## Run the Order Processing Example
 
 - Check out the [Durable Task Scheduler

durabletask-azuremanaged/CHANGELOG.md

@@ -7,6 +7,23 @@ adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
 ## Unreleased
 
+ADDED
+
+- Added optional `resource_id` configuration for the token audience on
+`DurableTaskSchedulerClient`, `AsyncDurableTaskSchedulerClient`,
+`DurableTaskSchedulerWorker`, `SandboxActivitiesClient`, and `SandboxWorker`.
+Explicit values override region defaults and support custom resource URIs.
+Surrounding whitespace, trailing slashes, and an existing `/.default` suffix
+are normalized before token requests; values that become empty are rejected.
+
+CHANGED
+
+- When `resource_id` is omitted or empty, Azure Managed clients and workers now
+use `https://durabletask.azure.us` if `REGION_NAME` starts with `usgov` or `usdod`
+(case-insensitive). Other regions, including an unset `REGION_NAME`, retain
+`https://durabletask.io`. The endpoint and credential authority remain separately
+configured.
+
 FIXED
 
 - With the corresponding core SDK update, asynchronous Azure Blob payload

durabletask-azuremanaged/durabletask/azuremanaged/client.py

@@ -25,10 +25,22 @@
 
 # Client class used for Durable Task Scheduler (DTS)
 class DurableTaskSchedulerClient(TaskHubGrpcClient):
+    """A client for Azure Durable Task Scheduler.
+
+    ``resource_id`` optionally overrides the token audience. If None or empty,
+    it defaults to ``https://durabletask.azure.us`` when ``REGION_NAME`` starts
+    with ``usgov`` or ``usdod`` (case-insensitive), or ``https://durabletask.io``
+    otherwise. Surrounding whitespace, trailing slashes, and an existing
+    ``/.default`` suffix are removed before requesting the ``/.default`` scope.
+    Values that become empty raise ``ValueError``. This does not configure the
+    service endpoint or the credential's authority.
+    """
+
     def __init__(self, *,
                  host_address: str,
                  taskhub: str,
                  token_credential: TokenCredential | None,
+                 resource_id: str | None = None,
                  channel: grpc.Channel | None = None,
                  secure_channel: bool = True,
                  interceptors: Sequence[shared.ClientInterceptor] | None = None,
@@ -47,7 +59,8 @@ def __init__(self, *,
         resolved_interceptors: list[shared.ClientInterceptor] = (
             list(interceptors) if interceptors is not None else []
         )
-        resolved_interceptors.append(DTSDefaultClientInterceptorImpl(token_credential, taskhub))
+        resolved_interceptors.append(DTSDefaultClientInterceptorImpl(
+            token_credential, taskhub, resource_id=resource_id))
 
         # We pass in None for the metadata so we don't construct an additional interceptor in the parent class
         # Since the parent class doesn't use anything metadata for anything else, we can set it as None
@@ -81,6 +94,12 @@ class AsyncDurableTaskSchedulerClient(AsyncTaskHubGrpcClient):
         taskhub (str): The name of the task hub. Cannot be empty.
         token_credential (TokenCredential | None): Azure credential for authentication.
             If None, anonymous authentication will be used.
+        resource_id (str | None, optional): Token audience override. If None or empty,
+            defaults to ``https://durabletask.azure.us`` when ``REGION_NAME`` starts
+            with ``usgov`` or ``usdod`` (case-insensitive), or ``https://durabletask.io``
+            otherwise. Surrounding whitespace, trailing slashes, and an existing
+            ``/.default`` suffix are removed before requesting the ``/.default`` scope.
+            Does not configure the service endpoint or the credential's authority.
         secure_channel (bool, optional): Whether to use a secure gRPC channel (TLS).
             Defaults to True.
         resiliency_options (GrpcClientResiliencyOptions | None, optional): Client-side
@@ -98,6 +117,7 @@ class AsyncDurableTaskSchedulerClient(AsyncTaskHubGrpcClient):
 
     Raises:
         ValueError: If taskhub is empty or None.
+        ValueError: If resource_id becomes empty after normalization.
 
     Example:
         >>> from azure.identity.aio import DefaultAzureCredential
@@ -116,6 +136,7 @@ def __init__(self, *,
                  host_address: str,
                  taskhub: str,
                  token_credential: AsyncTokenCredential | None,
+                 resource_id: str | None = None,
                  channel: grpc.aio.Channel | None = None,
                  secure_channel: bool = True,
                  interceptors: Sequence[shared.AsyncClientInterceptor] | None = None,
@@ -134,7 +155,8 @@ def __init__(self, *,
         resolved_interceptors: list[shared.AsyncClientInterceptor] = (
             list(interceptors) if interceptors is not None else []
         )
-        resolved_interceptors.append(DTSAsyncDefaultClientInterceptorImpl(token_credential, taskhub))
+        resolved_interceptors.append(DTSAsyncDefaultClientInterceptorImpl(
+            token_credential, taskhub, resource_id=resource_id))
 
         # We pass in None for the metadata so we don't construct an additional interceptor in the parent class
         # Since the parent class doesn't use anything metadata for anything else, we can set it as None

durabletask-azuremanaged/durabletask/azuremanaged/internal/access_token_manager.py

@@ -1,6 +1,7 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT License.
 import asyncio
+import os
 from datetime import datetime, timedelta, timezone
 from threading import Lock
 
@@ -10,14 +11,31 @@
 import durabletask.internal.shared as shared
 
 
+def resolve_resource_id(resource_id: str | None) -> str:
+    """Normalize an explicit token audience or resolve the current region's default."""
+    if resource_id is None or resource_id == "":
+        region = os.getenv("REGION_NAME", "")
+        if region.lower().startswith(("usgov", "usdod")):
+            return "https://durabletask.azure.us"
+        return "https://durabletask.io"
+
+    resource_id = resource_id.strip().rstrip("/")
+    if resource_id.lower().endswith("/.default"):
+        resource_id = resource_id[:-len("/.default")].rstrip("/")
+    if not resource_id:
+        raise ValueError("resource_id cannot be empty after normalization.")
+    return resource_id
+
+
 # By default, when there's 10minutes left before the token expires, refresh the token
 class AccessTokenManager:
 
     _token: AccessToken | None
     expiry_time: datetime | None
 
-    def __init__(self, token_credential: TokenCredential | None, refresh_interval_seconds: int = 600):
-        self._scope = "https://durabletask.io/.default"
+    def __init__(self, token_credential: TokenCredential | None, refresh_interval_seconds: int = 600,
+                 *, resource_id: str | None = None):
+        self._scope = f"{resolve_resource_id(resource_id)}/.default"
         self._refresh_interval_seconds = refresh_interval_seconds
         self._logger = shared.get_logger("token_manager")
 
@@ -63,8 +81,8 @@ class AsyncAccessTokenManager:
     _token: AccessToken | None
 
     def __init__(self, token_credential: AsyncTokenCredential | None,
-                 refresh_interval_seconds: int = 600):
-        self._scope = "https://durabletask.io/.default"
+                 refresh_interval_seconds: int = 600, *, resource_id: str | None = None):
+        self._scope = f"{resolve_resource_id(resource_id)}/.default"
         self._refresh_interval_seconds = refresh_interval_seconds
         self._logger = shared.get_logger("async_token_manager")
 

durabletask-azuremanaged/durabletask/azuremanaged/internal/durabletask_grpc_interceptor.py

@@ -11,6 +11,7 @@
 from durabletask.azuremanaged.internal.access_token_manager import (
     AccessTokenManager,
     AsyncAccessTokenManager,
+    resolve_resource_id,
 )
 from durabletask.internal.grpc_interceptor import (
     DefaultAsyncClientInterceptorImpl,
@@ -43,7 +44,10 @@ def __init__(
             self,
             token_credential: TokenCredential | None,
             taskhub_name: str,
-            worker_id: str | None = None):
+            worker_id: str | None = None,
+            *, resource_id: str | None = None):
+        if token_credential is None:
+            resolve_resource_id(resource_id)
         user_agent = f"durabletask-python/{_get_sdk_version()}"
         self._metadata = [
             ("taskhub", taskhub_name),
@@ -58,7 +62,8 @@ def __init__(
         self._token_manager = None
         if token_credential is not None:
             self._token_credential = token_credential
-            self._token_manager = AccessTokenManager(token_credential=self._token_credential)
+            self._token_manager = AccessTokenManager(
+                token_credential=self._token_credential, resource_id=resource_id)
 
     def _upsert_authorization_header(self, token: str) -> None:
         found = False
@@ -91,7 +96,10 @@ class DTSAsyncDefaultClientInterceptorImpl(DefaultAsyncClientInterceptorImpl):
     This class implements async gRPC interceptors to add DTS-specific headers
     (task hub name, user agent, and authentication token) to all async calls."""
 
-    def __init__(self, token_credential: AsyncTokenCredential | None, taskhub_name: str):
+    def __init__(self, token_credential: AsyncTokenCredential | None, taskhub_name: str,
+                 *, resource_id: str | None = None):
+        if token_credential is None:
+            resolve_resource_id(resource_id)
         user_agent = f"durabletask-python/{_get_sdk_version()}"
         self._metadata = [
             ("taskhub", taskhub_name),
@@ -104,7 +112,8 @@ def __init__(self, token_credential: AsyncTokenCredential | None, taskhub_name:
         self._token_manager = None
         if token_credential is not None:
             self._token_credential = token_credential
-            self._token_manager = AsyncAccessTokenManager(token_credential=self._token_credential)
+            self._token_manager = AsyncAccessTokenManager(
+                token_credential=self._token_credential, resource_id=resource_id)
 
     def _upsert_authorization_header(self, token: str) -> None:
         found = False

durabletask-azuremanaged/durabletask/azuremanaged/preview/sandboxes/client.py

@@ -18,13 +18,19 @@
 
 
 class SandboxActivitiesClient:
-    """Client for Durable Task Scheduler sandbox activity management operations."""
+    """Client for Durable Task Scheduler sandbox activity management operations.
+
+    ``resource_id`` selects the token audience using the same normalization and
+    ``REGION_NAME`` defaults as ``DurableTaskSchedulerClient``. It does not change
+    the endpoint or credential authority. A supplied channel owns its authentication.
+    """
 
     def __init__(
             self, *,
             host_address: str,
             taskhub: str,
             token_credential: Optional[TokenCredential],
+            resource_id: Optional[str] = None,
             channel: Optional[grpc.Channel] = None,
             secure_channel: bool = True,
             interceptors: Optional[Sequence[shared.ClientInterceptor]] = None,
@@ -33,6 +39,7 @@ def __init__(
             host_address=host_address,
             taskhub=taskhub,
             token_credential=token_credential,
+            resource_id=resource_id,
             channel=channel,
             secure_channel=secure_channel,
             interceptors=interceptors,

durabletask-azuremanaged/durabletask/azuremanaged/preview/sandboxes/transport.py

@@ -40,6 +40,7 @@ def __init__(
             host_address: str,
             taskhub: str,
             token_credential: Optional[TokenCredential],
+            resource_id: Optional[str] = None,
             channel: Optional[grpc.Channel] = None,
             secure_channel: bool = True,
             interceptors: Optional[Sequence[shared.ClientInterceptor]] = None,
@@ -52,7 +53,8 @@ def __init__(
             resolved_interceptors: list[shared.ClientInterceptor] = (
                 list(interceptors) if interceptors is not None else []
             )
-            resolved_interceptors.append(DTSDefaultClientInterceptorImpl(token_credential, taskhub))
+            resolved_interceptors.append(DTSDefaultClientInterceptorImpl(
+                token_credential, taskhub, resource_id=resource_id))
             channel = shared.get_grpc_channel(
                 host_address=host_address,
                 secure_channel=secure_channel,

durabletask-azuremanaged/durabletask/azuremanaged/preview/sandboxes/worker.py

@@ -12,6 +12,7 @@
 from azure.identity import ManagedIdentityCredential
 
 from durabletask.azuremanaged.internal import sandbox_service_pb2 as pb
+from durabletask.azuremanaged.internal.access_token_manager import resolve_resource_id
 from durabletask.azuremanaged.preview.sandboxes.helpers import SandboxActivity
 from durabletask.azuremanaged.preview.sandboxes.helpers import resolve_activities
 from durabletask.azuremanaged.preview.sandboxes.worker_profiles import (
@@ -38,9 +39,15 @@ class SandboxWorker(DurableTaskSchedulerWorker):
 
     This worker registers a live worker session with Durable Task Scheduler and
     restricts dispatch to the activities registered on this worker.
+
+    ``resource_id`` selects the token audience using the same normalization and
+    ``REGION_NAME`` defaults as ``DurableTaskSchedulerWorker``. It applies to both
+    activity execution and worker registration without changing the runtime endpoint
+    or managed identity configuration.
     """
 
-    def __init__(self) -> None:
+    def __init__(self, *, resource_id: str | None = None) -> None:
+        resolved_resource_id = resource_id or resolve_resource_id(None)
         resolved_host_address = _resolve_host_address()
         resolved_taskhub = _resolve_taskhub()
         resolved_secure_channel = _resolve_secure_channel(resolved_host_address)
@@ -54,12 +61,14 @@ def __init__(self) -> None:
         self._sandbox_host_address = resolved_host_address
         self._sandbox_secure_channel = resolved_secure_channel
         self._sandbox_token_credential = resolved_token_credential
+        self._sandbox_resource_id = resolved_resource_id
         self._sandbox_logger = shared.get_logger("worker")
 
         super().__init__(
             host_address=resolved_host_address,
             taskhub=resolved_taskhub,
             token_credential=resolved_token_credential,
+            resource_id=resolved_resource_id,
             secure_channel=resolved_secure_channel,
             concurrency_options=concurrency_options)
 
@@ -139,6 +148,7 @@ def _run_sandbox_registration_loop(self) -> None:
                             host_address=self._sandbox_host_address,
                             taskhub=self._sandbox_taskhub,
                             token_credential=self._sandbox_token_credential,
+                            resource_id=self._sandbox_resource_id,
                             secure_channel=self._sandbox_secure_channel)
                     client.connect_sandbox_activity_worker(self._registration_messages())
                     retry_delay = 1.0