py2docfx 0.1.20rc2196756__py3-none-any.whl → 0.1.21.dev2246704__py3-none-any.whl

py2docfx/venv/venv1/Lib/site-packages/azure/identity/aio/_credentials/azd_cli.py

@@ -21,10 +21,12 @@ from ..._credentials.azd_cli import (
     EXECUTABLE_NAME,
     get_safe_working_dir,
     NOT_LOGGED_IN,
+    UNKNOWN_CLAIMS_FLAG,
     parse_token,
     sanitize_output,
+    extract_cli_error_message,
 )
-from ..._internal import resolve_tenant, within_dac, validate_tenant_id, validate_scope
+from ..._internal import encode_base64, resolve_tenant, within_dac, validate_tenant_id, validate_scope
 
 
 _LOGGER = logging.getLogger(__name__)
@@ -87,7 +89,7 @@ class AzureDeveloperCliCredential(AsyncContextManager):
     async def get_token(
         self,
         *scopes: str,
-        claims: Optional[str] = None,  # pylint:disable=unused-argument
+        claims: Optional[str] = None,
         tenant_id: Optional[str] = None,
         **kwargs: Any,
     ) -> AccessToken:
@@ -99,7 +101,8 @@ class AzureDeveloperCliCredential(AsyncContextManager):
         :param str scopes: desired scope for the access token. This credential allows only one scope per request.
             For more information about scopes, see
             https://learn.microsoft.com/entra/identity-platform/scopes-oidc.
-        :keyword str claims: not used by this credential; any value provided will be ignored.
+        :keyword str claims: additional claims required in the token, such as those returned in a resource provider's
+            claims challenge following an authorization failure.
         :keyword str tenant_id: optional tenant to include in the token request.
 
         :return: An access token with the desired scopes.
@@ -110,11 +113,13 @@ class AzureDeveloperCliCredential(AsyncContextManager):
         """
         # only ProactorEventLoop supports subprocesses on Windows (and it isn't the default loop on Python < 3.8)
         if sys.platform.startswith("win") and not isinstance(asyncio.get_event_loop(), asyncio.ProactorEventLoop):
-            return _SyncAzureDeveloperCliCredential().get_token(*scopes, tenant_id=tenant_id, **kwargs)
+            return _SyncAzureDeveloperCliCredential().get_token(*scopes, claims=claims, tenant_id=tenant_id, **kwargs)
 
         options: TokenRequestOptions = {}
         if tenant_id:
             options["tenant_id"] = tenant_id
+        if claims:
+            options["claims"] = claims
 
         token_info = await self._get_token_base(*scopes, options=options, **kwargs)
         return AccessToken(token_info.token, token_info.expires_on)
@@ -152,6 +157,7 @@ class AzureDeveloperCliCredential(AsyncContextManager):
             raise ValueError("Missing scope in request. \n")
 
         tenant_id = options.get("tenant_id") if options else None
+        claims = options.get("claims") if options else None
         if tenant_id:
             validate_tenant_id(tenant_id)
         for scope in scopes:
@@ -169,16 +175,22 @@ class AzureDeveloperCliCredential(AsyncContextManager):
 
         if tenant:
             command_args += ["--tenant-id", tenant]
+        if claims:
+            command_args += ["--claims", encode_base64(claims)]
         output = await _run_command(command_args, self._process_timeout)
 
         token = parse_token(output)
         if not token:
-            sanitized_output = sanitize_output(output)
-            message = (
-                f"Unexpected output from Azure Developer CLI: '{sanitized_output}'. \n"
-                f"To mitigate this issue, please refer to the troubleshooting guidelines here at "
-                f"https://aka.ms/azsdk/python/identity/azdevclicredential/troubleshoot."
-            )
+            extracted = extract_cli_error_message(output)
+            if extracted:
+                message = extracted
+            else:
+                sanitized_output = sanitize_output(output)
+                message = (
+                    f"Unexpected output from Azure Developer CLI: '{sanitized_output}'. \n"
+                    f"To mitigate this issue, please refer to the troubleshooting guidelines here at "
+                    f"https://aka.ms/azsdk/python/identity/azdevclicredential/troubleshoot."
+                )
             if within_dac.get():
                 raise CredentialUnavailableError(message=message)
             raise ClientAuthenticationError(message=message)
@@ -226,10 +238,17 @@ async def _run_command(command_args: List[str], timeout: int) -> str:
     if proc.returncode == 127 or stderr.startswith("'azd' is not recognized"):
         raise CredentialUnavailableError(CLI_NOT_FOUND)
 
-    if "not logged in, run `azd auth login` to login" in stderr and "AADSTS" not in stderr:
+    combined_text = f"{output}\n{stderr}"
+    if "not logged in, run `azd auth login` to login" in combined_text and "AADSTS" not in combined_text:
         raise CredentialUnavailableError(message=NOT_LOGGED_IN)
-
-    message = sanitize_output(stderr) if stderr else "Failed to invoke Azure Developer CLI"
+    if "unknown flag: --claims" in combined_text:
+        raise CredentialUnavailableError(message=UNKNOWN_CLAIMS_FLAG)
+
+    message = (
+        extract_cli_error_message(output)
+        or extract_cli_error_message(stderr)
+        or (sanitize_output(stderr) if stderr else "Failed to invoke Azure Developer CLI")
+    )
     if within_dac.get():
         raise CredentialUnavailableError(message=message)
     raise ClientAuthenticationError(message=message)

py2docfx/venv/venv1/Lib/site-packages/azure/identity/aio/_credentials/azure_cli.py

@@ -23,9 +23,11 @@ from ..._credentials.azure_cli import (
     NOT_LOGGED_IN,
     parse_token,
     sanitize_output,
+    CLAIMS_UNSUPPORTED_ERROR,
 )
 from ..._internal import (
     _scopes_to_resource,
+    encode_base64,
     resolve_tenant,
     within_dac,
     validate_tenant_id,
@@ -82,7 +84,7 @@ class AzureCliCredential(AsyncContextManager):
     async def get_token(
         self,
         *scopes: str,
-        claims: Optional[str] = None,  # pylint:disable=unused-argument
+        claims: Optional[str] = None,
         tenant_id: Optional[str] = None,
         **kwargs: Any,
     ) -> AccessToken:
@@ -94,22 +96,26 @@ class AzureCliCredential(AsyncContextManager):
         :param str scopes: desired scope for the access token. This credential allows only one scope per request.
             For more information about scopes, see
             https://learn.microsoft.com/entra/identity-platform/scopes-oidc.
-        :keyword str claims: not used by this credential; any value provided will be ignored.
+        :keyword str claims: additional claims required in the token. This credential does not support claims
+            challenges.
         :keyword str tenant_id: optional tenant to include in the token request.
 
         :return: An access token with the desired scopes.
         :rtype: ~azure.core.credentials.AccessToken
-        :raises ~azure.identity.CredentialUnavailableError: the credential was unable to invoke the Azure CLI.
+        :raises ~azure.identity.CredentialUnavailableError: the credential was either unable to invoke the Azure CLI
+          or a claims challenge was provided.
         :raises ~azure.core.exceptions.ClientAuthenticationError: the credential invoked the Azure CLI but didn't
           receive an access token.
         """
         # only ProactorEventLoop supports subprocesses on Windows (and it isn't the default loop on Python < 3.8)
         if sys.platform.startswith("win") and not isinstance(asyncio.get_event_loop(), asyncio.ProactorEventLoop):
-            return _SyncAzureCliCredential().get_token(*scopes, tenant_id=tenant_id, **kwargs)
+            return _SyncAzureCliCredential().get_token(*scopes, claims=claims, tenant_id=tenant_id, **kwargs)
 
         options: TokenRequestOptions = {}
         if tenant_id:
             options["tenant_id"] = tenant_id
+        if claims:
+            options["claims"] = claims
 
         token_info = await self._get_token_base(*scopes, options=options, **kwargs)
         return AccessToken(token_info.token, token_info.expires_on)
@@ -130,7 +136,8 @@ class AzureCliCredential(AsyncContextManager):
         :rtype: ~azure.core.credentials.AccessTokenInfo
         :return: An AccessTokenInfo instance containing information about the token.
 
-        :raises ~azure.identity.CredentialUnavailableError: the credential was unable to invoke the Azure CLI.
+        :raises ~azure.identity.CredentialUnavailableError: the credential was either unable to invoke the Azure CLI
+          or a claims challenge was provided.
         :raises ~azure.core.exceptions.ClientAuthenticationError: the credential invoked the Azure CLI but didn't
           receive an access token.
         """
@@ -142,6 +149,20 @@ class AzureCliCredential(AsyncContextManager):
     async def _get_token_base(
         self, *scopes: str, options: Optional[TokenRequestOptions] = None, **kwargs: Any
     ) -> AccessTokenInfo:
+        # Check for claims challenge first
+        if options and "claims" in options and options["claims"]:
+            error_message = CLAIMS_UNSUPPORTED_ERROR.format(claims_value=encode_base64(options["claims"]))
+
+            # Add tenant if provided in options
+            if options.get("tenant_id"):
+                error_message += f" --tenant {options.get('tenant_id')}"
+
+            # Add scope if provided
+            if scopes:
+                error_message += f" --scope {scopes[0]}"
+
+            raise CredentialUnavailableError(message=error_message)
+
         tenant_id = options.get("tenant_id") if options else None
         if tenant_id:
             validate_tenant_id(tenant_id)

py2docfx/venv/venv1/Lib/site-packages/azure/identity/aio/_credentials/azure_powershell.py

@@ -16,8 +16,9 @@ from ..._credentials.azure_powershell import (
     get_safe_working_dir,
     raise_for_error,
     parse_token,
+    CLAIMS_UNSUPPORTED_ERROR,
 )
-from ..._internal import resolve_tenant, validate_tenant_id, validate_scope
+from ..._internal import encode_base64, resolve_tenant, validate_tenant_id, validate_scope
 
 
 class AzurePowerShellCredential(AsyncContextManager):
@@ -58,7 +59,7 @@ class AzurePowerShellCredential(AsyncContextManager):
     async def get_token(
         self,
         *scopes: str,
-        claims: Optional[str] = None,  # pylint:disable=unused-argument
+        claims: Optional[str] = None,
         tenant_id: Optional[str] = None,
         **kwargs: Any,
     ) -> AccessToken:
@@ -87,6 +88,8 @@ class AzurePowerShellCredential(AsyncContextManager):
         options: TokenRequestOptions = {}
         if tenant_id:
             options["tenant_id"] = tenant_id
+        if claims:
+            options["claims"] = claims
 
         token_info = await self._get_token_base(*scopes, options=options, **kwargs)
         return AccessToken(token_info.token, token_info.expires_on)
@@ -119,6 +122,14 @@ class AzurePowerShellCredential(AsyncContextManager):
     async def _get_token_base(
         self, *scopes: str, options: Optional[TokenRequestOptions] = None, **kwargs: Any
     ) -> AccessTokenInfo:
+
+        # Check if claims challenge is provided
+        if options and "claims" in options and options["claims"]:
+            error_message = CLAIMS_UNSUPPORTED_ERROR.format(claims_value=encode_base64(options["claims"]))
+            if options.get("tenant_id"):
+                error_message += f" -Tenant {options.get('tenant_id')}"
+            raise CredentialUnavailableError(message=error_message)
+
         tenant_id = options.get("tenant_id") if options else None
         if tenant_id:
             validate_tenant_id(tenant_id)

py2docfx/venv/venv1/Lib/site-packages/azure/identity/aio/_credentials/chained.py

@@ -25,7 +25,7 @@ class ChainedTokenCredential(AsyncContextManager):
     https://aka.ms/azsdk/python/identity/credential-chains#chainedtokencredential-overview.
 
     :param credentials: credential instances to form the chain
-    :type credentials: ~azure.core.credentials_async.AsyncTokenCredential
+    :type credentials: ~azure.core.credentials_async.AsyncTokenProvider
 
     .. admonition:: Example:
 

py2docfx/venv/venv1/Lib/site-packages/azure/identity/aio/_credentials/default.py

@@ -8,8 +8,9 @@ from typing import List, Optional, Any, cast
 
 from azure.core.credentials import AccessToken, AccessTokenInfo, TokenRequestOptions
 from azure.core.credentials_async import AsyncTokenCredential, AsyncSupportsTokenInfo
+from ... import CredentialUnavailableError
 from ..._constants import EnvironmentVariables
-from ..._internal import get_default_authority, normalize_authority, within_dac
+from ..._internal import get_default_authority, normalize_authority, within_dac, process_credential_exclusions
 from .azure_cli import AzureCliCredential
 from .azd_cli import AzureDeveloperCliCredential
 from .azure_powershell import AzurePowerShellCredential
@@ -24,6 +25,34 @@ from .workload_identity import WorkloadIdentityCredential
 _LOGGER = logging.getLogger(__name__)
 
 
+class AsyncFailedDACCredential:
+    """Async version of FailedDACCredential for use in async credential chains.
+
+    This acts as a substitute for an async credential that has failed to initialize in the DAC chain.
+    """
+
+    def __init__(self, credential_name: str, error: str) -> None:
+        self._error = error
+        self._credential_name = credential_name
+
+    async def get_token(self, *scopes: str, **kwargs: Any) -> AccessToken:
+        raise CredentialUnavailableError(self._error)
+
+    async def get_token_info(
+        self, *scopes, options: Optional[TokenRequestOptions] = None, **kwargs: Any
+    ) -> AccessTokenInfo:
+        raise CredentialUnavailableError(self._error)
+
+    async def __aexit__(self, *args: Any, **kwargs: Any) -> None:
+        pass
+
+    async def __aenter__(self) -> "AsyncFailedDACCredential":
+        return self
+
+    async def close(self) -> None:
+        pass
+
+
 class DefaultAzureCredential(ChainedTokenCredential):
     """A credential capable of handling most Azure SDK authentication scenarios. See
     https://aka.ms/azsdk/python/identity/credential-chains#usage-guidance-for-defaultazurecredential.
@@ -73,11 +102,13 @@ class DefaultAzureCredential(ChainedTokenCredential):
     :keyword str shared_cache_tenant_id: Preferred tenant for :class:`~azure.identity.aio.SharedTokenCacheCredential`.
         Defaults to the value of environment variable AZURE_TENANT_ID, if any.
     :keyword str visual_studio_code_tenant_id: Tenant ID to use when authenticating with
-        :class:`~azure.identity.aio.VisualStudioCodeCredential`. Defaults to the "Azure: Tenant" setting in VS Code's
-        user settings or, when that setting has no value, the "organizations" tenant, which supports only Azure Active
-        Directory work or school accounts.
+        :class:`~azure.identity.VisualStudioCodeCredential`. Defaults to the tenant specified in the authentication
+        record file used by the Azure Resources extension.
     :keyword int process_timeout: The timeout in seconds to use for developer credentials that run
         subprocesses (e.g. AzureCliCredential, AzurePowerShellCredential). Defaults to **10** seconds.
+    :keyword bool require_envvar: If **True**, require that the AZURE_TOKEN_CREDENTIALS environment variable be set
+        to a value denoting the credential type or credential group to use. If unset or empty, DefaultAzureCredential
+        will raise a `ValueError`. Defaults to **False**.
 
     .. admonition:: Example:
 
@@ -89,23 +120,15 @@ class DefaultAzureCredential(ChainedTokenCredential):
             :caption: Create a DefaultAzureCredential.
     """
 
-    def __init__(self, **kwargs: Any) -> None:  # pylint: disable=too-many-statements
+    def __init__(self, **kwargs: Any) -> None:  # pylint: disable=too-many-statements, too-many-locals
         if "tenant_id" in kwargs:
             raise TypeError("'tenant_id' is not supported in DefaultAzureCredential.")
 
         authority = kwargs.pop("authority", None)
-
-        vscode_tenant_id = kwargs.pop(
-            "visual_studio_code_tenant_id", os.environ.get(EnvironmentVariables.AZURE_TENANT_ID)
-        )
-        vscode_args = dict(kwargs)
-        if authority:
-            vscode_args["authority"] = authority
-        if vscode_tenant_id:
-            vscode_args["tenant_id"] = vscode_tenant_id
-
         authority = normalize_authority(authority) if authority else get_default_authority()
 
+        vscode_tenant_id = kwargs.pop("visual_studio_code_tenant_id", None)
+
         shared_cache_username = kwargs.pop("shared_cache_username", os.environ.get(EnvironmentVariables.AZURE_USERNAME))
         shared_cache_tenant_id = kwargs.pop(
             "shared_cache_tenant_id", os.environ.get(EnvironmentVariables.AZURE_TENANT_ID)
@@ -119,56 +142,94 @@ class DefaultAzureCredential(ChainedTokenCredential):
             "workload_identity_tenant_id", os.environ.get(EnvironmentVariables.AZURE_TENANT_ID)
         )
 
-        vscode_tenant_id = kwargs.pop(
-            "visual_studio_code_tenant_id", os.environ.get(EnvironmentVariables.AZURE_TENANT_ID)
-        )
-
         process_timeout = kwargs.pop("process_timeout", 10)
-
-        token_credentials_env = os.environ.get(EnvironmentVariables.AZURE_TOKEN_CREDENTIALS, "").strip().lower()
-        exclude_workload_identity_credential = kwargs.pop("exclude_workload_identity_credential", False)
-        exclude_visual_studio_code_credential = kwargs.pop("exclude_visual_studio_code_credential", True)
-        exclude_developer_cli_credential = kwargs.pop("exclude_developer_cli_credential", False)
-        exclude_cli_credential = kwargs.pop("exclude_cli_credential", False)
-        exclude_environment_credential = kwargs.pop("exclude_environment_credential", False)
-        exclude_managed_identity_credential = kwargs.pop("exclude_managed_identity_credential", False)
-        exclude_shared_token_cache_credential = kwargs.pop("exclude_shared_token_cache_credential", False)
-        exclude_powershell_credential = kwargs.pop("exclude_powershell_credential", False)
-
-        if token_credentials_env == "dev":
-            # In dev mode, use only developer credentials
-            exclude_environment_credential = True
-            exclude_managed_identity_credential = True
-            exclude_workload_identity_credential = True
-        elif token_credentials_env == "prod":
-            # In prod mode, use only production credentials
-            exclude_shared_token_cache_credential = True
-            exclude_visual_studio_code_credential = True
-            exclude_cli_credential = True
-            exclude_developer_cli_credential = True
-            exclude_powershell_credential = True
-        elif token_credentials_env != "":
-            # If the environment variable is set to something other than dev or prod, raise an error
+        require_envvar = kwargs.pop("require_envvar", False)
+        if require_envvar and not os.environ.get(EnvironmentVariables.AZURE_TOKEN_CREDENTIALS):
             raise ValueError(
-                f"Invalid value for {EnvironmentVariables.AZURE_TOKEN_CREDENTIALS}: {token_credentials_env}. "
-                "Valid values are 'dev' or 'prod'."
+                "AZURE_TOKEN_CREDENTIALS environment variable is required but is not set or is empty. "
+                "Set it to 'dev', 'prod', or a specific credential name."
             )
 
+        # Define credential configuration mapping (async version)
+        credential_config = {
+            "environment": {
+                "exclude_param": "exclude_environment_credential",
+                "env_name": "environmentcredential",
+                "default_exclude": False,
+            },
+            "workload_identity": {
+                "exclude_param": "exclude_workload_identity_credential",
+                "env_name": "workloadidentitycredential",
+                "default_exclude": False,
+            },
+            "managed_identity": {
+                "exclude_param": "exclude_managed_identity_credential",
+                "env_name": "managedidentitycredential",
+                "default_exclude": False,
+            },
+            "shared_token_cache": {
+                "exclude_param": "exclude_shared_token_cache_credential",
+                "default_exclude": False,
+            },
+            "visual_studio_code": {
+                "exclude_param": "exclude_visual_studio_code_credential",
+                "env_name": "visualstudiocodecredential",
+                "default_exclude": False,
+            },
+            "cli": {
+                "exclude_param": "exclude_cli_credential",
+                "env_name": "azureclicredential",
+                "default_exclude": False,
+            },
+            "developer_cli": {
+                "exclude_param": "exclude_developer_cli_credential",
+                "env_name": "azuredeveloperclicredential",
+                "default_exclude": False,
+            },
+            "powershell": {
+                "exclude_param": "exclude_powershell_credential",
+                "env_name": "azurepowershellcredential",
+                "default_exclude": False,
+            },
+        }
+
+        # Extract user-provided exclude flags and set defaults
+        exclude_flags = {}
+        user_excludes = {}
+        for cred_key, config in credential_config.items():
+            param_name = cast(str, config["exclude_param"])
+            user_excludes[cred_key] = kwargs.pop(param_name, None)
+            exclude_flags[cred_key] = config["default_exclude"]
+
+        # Process AZURE_TOKEN_CREDENTIALS environment variable and apply user overrides
+        exclude_flags = process_credential_exclusions(credential_config, exclude_flags, user_excludes)
+
+        # Extract individual exclude flags for backward compatibility
+        exclude_environment_credential = exclude_flags["environment"]
+        exclude_workload_identity_credential = exclude_flags["workload_identity"]
+        exclude_managed_identity_credential = exclude_flags["managed_identity"]
+        exclude_shared_token_cache_credential = exclude_flags["shared_token_cache"]
+        exclude_visual_studio_code_credential = exclude_flags["visual_studio_code"]
+        exclude_cli_credential = exclude_flags["cli"]
+        exclude_developer_cli_credential = exclude_flags["developer_cli"]
+        exclude_powershell_credential = exclude_flags["powershell"]
+
         credentials: List[AsyncSupportsTokenInfo] = []
         within_dac.set(True)
         if not exclude_environment_credential:
             credentials.append(EnvironmentCredential(authority=authority, _within_dac=True, **kwargs))
         if not exclude_workload_identity_credential:
-            if all(os.environ.get(var) for var in EnvironmentVariables.WORKLOAD_IDENTITY_VARS):
-                client_id = workload_identity_client_id
+            try:
                 credentials.append(
                     WorkloadIdentityCredential(
-                        client_id=cast(str, client_id),
+                        client_id=cast(str, workload_identity_client_id),
                         tenant_id=workload_identity_tenant_id,
-                        token_file_path=os.environ[EnvironmentVariables.AZURE_FEDERATED_TOKEN_FILE],
+                        token_file_path=os.environ.get(EnvironmentVariables.AZURE_FEDERATED_TOKEN_FILE),
                         **kwargs,
                     )
                 )
+            except ValueError as ex:
+                credentials.append(AsyncFailedDACCredential("WorkloadIdentityCredential", error=str(ex)))
         if not exclude_managed_identity_credential:
             credentials.append(
                 ManagedIdentityCredential(
@@ -187,7 +248,7 @@ class DefaultAzureCredential(ChainedTokenCredential):
             except Exception as ex:  # pylint:disable=broad-except
                 _LOGGER.info("Shared token cache is unavailable: '%s'", ex)
         if not exclude_visual_studio_code_credential:
-            credentials.append(VisualStudioCodeCredential(**vscode_args))
+            credentials.append(VisualStudioCodeCredential(tenant_id=vscode_tenant_id))
         if not exclude_cli_credential:
             credentials.append(AzureCliCredential(process_timeout=process_timeout))
         if not exclude_powershell_credential:
@@ -226,8 +287,10 @@ class DefaultAzureCredential(ChainedTokenCredential):
             return token
 
         within_dac.set(True)
-        token = await super().get_token(*scopes, claims=claims, tenant_id=tenant_id, **kwargs)
-        within_dac.set(False)
+        try:
+            token = await super().get_token(*scopes, claims=claims, tenant_id=tenant_id, **kwargs)
+        finally:
+            within_dac.set(False)
         return token
 
     async def get_token_info(self, *scopes: str, options: Optional[TokenRequestOptions] = None) -> AccessTokenInfo:
@@ -257,6 +320,8 @@ class DefaultAzureCredential(ChainedTokenCredential):
             return token_info
 
         within_dac.set(True)
-        token_info = await cast(AsyncSupportsTokenInfo, super()).get_token_info(*scopes, options=options)
-        within_dac.set(False)
+        try:
+            token_info = await cast(AsyncSupportsTokenInfo, super()).get_token_info(*scopes, options=options)
+        finally:
+            within_dac.set(False)
         return token_info

py2docfx/venv/venv1/Lib/site-packages/azure/identity/aio/_credentials/imds.py

@@ -3,10 +3,13 @@
 # Licensed under the MIT License.
 # ------------------------------------
 import os
-from typing import Optional, Any
+from typing import Optional, Any, Dict
 
 from azure.core.exceptions import ClientAuthenticationError, HttpResponseError
 from azure.core.credentials import AccessTokenInfo
+from azure.core.pipeline.policies import AsyncRetryPolicy
+from azure.core.pipeline import PipelineResponse
+
 from ... import CredentialUnavailableError
 from ..._constants import EnvironmentVariables
 from .._internal import AsyncContextManager
@@ -16,10 +19,33 @@ from ..._internal import within_credential_chain
 from ..._credentials.imds import _get_request, _check_forbidden_response, PIPELINE_SETTINGS
 
 
+class AsyncImdsRetryPolicy(AsyncRetryPolicy):
+    """Async custom retry policy for IMDS credential with extended retry duration for 410 responses.
+
+    This policy ensures that specifically for 410 status codes, the total exponential backoff duration
+    is at least 70 seconds to handle temporary IMDS endpoint unavailability.
+    For other status codes, it uses the standard retry behavior.
+    """
+
+    def __init__(self, **kwargs: Any) -> None:
+        # Increased backoff factor to ensure at least 70 seconds retry duration for 410 responses.
+        # Five retries, with each retry sleeping for [0.0s, 5.0s, 10.0s, 20.0s, 40.0s] between attempts (75s total)
+        self.backoff_factor_for_410 = 2.5
+        super().__init__(**kwargs)
+
+    def is_retry(self, settings: Dict[str, Any], response: PipelineResponse[Any, Any]) -> bool:
+        if response.http_response.status_code == 410:
+            settings["backoff"] = self.backoff_factor_for_410
+        else:
+            settings["backoff"] = self.backoff_factor
+        return super().is_retry(settings, response)
+
+
 class ImdsCredential(AsyncContextManager, GetTokenMixin):
     def __init__(self, **kwargs: Any) -> None:
         super().__init__()
 
+        kwargs["retry_policy_class"] = AsyncImdsRetryPolicy
         self._client = AsyncManagedIdentityClient(_get_request, **dict(PIPELINE_SETTINGS, **kwargs))
         if EnvironmentVariables.AZURE_POD_IDENTITY_AUTHORITY_HOST in os.environ:
             self._endpoint_available: Optional[bool] = True

py2docfx/venv/venv1/Lib/site-packages/azure/identity/aio/_credentials/on_behalf_of.py

@@ -119,7 +119,7 @@ class OnBehalfOfCredential(AsyncContextManager, GetTokenMixin):
         # Note we assume the cache has tokens for one user only. That's okay because each instance of this class is
         # locked to a single user (assertion). This assumption will become unsafe if this class allows applications
         # to change an instance's assertion.
-        refresh_tokens = self._client.get_cached_refresh_tokens(scopes)
+        refresh_tokens = self._client.get_cached_refresh_tokens(scopes, **kwargs)
         if len(refresh_tokens) == 1:  # there should be only one
             try:
                 refresh_token = refresh_tokens[0]["secret"]

py2docfx/venv/venv1/Lib/site-packages/azure/identity/aio/_credentials/shared_cache.py

@@ -8,6 +8,7 @@ from ..._internal.aad_client import AadClientBase
 from ... import CredentialUnavailableError
 from ..._constants import DEVELOPER_SIGN_ON_CLIENT_ID
 from ..._internal.shared_token_cache import NO_TOKEN, SharedTokenCacheBase
+from ..._internal.utils import within_dac
 from .._internal import AsyncContextManager
 from .._internal.aad_client import AadClient
 from .._internal.decorators import log_get_token_async
@@ -143,11 +144,17 @@ class SharedTokenCacheCredential(SharedTokenCacheBase, AsyncContextManager):
 
         # try each refresh token, returning the first access token acquired
         for refresh_token in self._get_refresh_tokens(account, is_cae=is_cae):
-            token = await cast(AadClient, self._client).obtain_token_by_refresh_token(
-                scopes, refresh_token, claims=claims, tenant_id=tenant_id, enable_cae=is_cae, **kwargs
-            )
-            return token
-
+            try:
+                token = await cast(AadClient, self._client).obtain_token_by_refresh_token(
+                    scopes, refresh_token, claims=claims, tenant_id=tenant_id, enable_cae=is_cae, **kwargs
+                )
+                return token
+            except Exception as e:  # pylint: disable=broad-except
+                if within_dac.get():
+                    raise CredentialUnavailableError(  # pylint: disable=raise-missing-from
+                        message=getattr(e, "message", str(e)), response=getattr(e, "response", None)
+                    )
+                raise
         raise CredentialUnavailableError(message=NO_TOKEN.format(account.get("username")))
 
     def _get_auth_client(self, **kwargs: Any) -> AadClientBase: