Python: Unify Azure credential handling across all packages#4087
Python: Unify Azure credential handling across all packages#4087eavanvalkenburg wants to merge 1 commit intomicrosoft:mainfrom
Conversation
Replace ad_token, ad_token_provider, and get_entra_auth_token with a unified credential parameter across all Azure-related packages. Core changes: - Add AzureCredentialTypes (TokenCredential | AsyncTokenCredential) and AzureTokenProvider (Callable[[], str | Awaitable[str]]) type aliases - Add resolve_credential_to_token_provider() using azure.identity's get_bearer_token_provider for automatic token caching/refresh - Update AzureOpenAIChatClient, AzureOpenAIResponsesClient, and AzureOpenAIAssistantsClient to accept credential: AzureCredentialTypes | AzureTokenProvider - Remove ad_token, ad_token_provider params and get_entra_auth_token helpers Package updates: - azure-ai: Accept AzureCredentialTypes on AzureAIClient, AzureAIAgentClient, AzureAIProjectAgentProvider, AzureAIAgentsProvider - azure-ai-search: Accept AzureCredentialTypes on AzureAISearchContextProvider - purview: Accept AzureCredentialTypes | AzureTokenProvider on PurviewClient, PurviewPolicyMiddleware, PurviewChatPolicyMiddleware Fixes microsoft#3449 Fixes microsoft#3500 Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Python Test Coverage Report •
Python Unit Test Overview
|
||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
There was a problem hiding this comment.
Pull request overview
This PR standardizes Microsoft Entra ID authentication across the Python Azure integrations by replacing the older ad_token / ad_token_provider / get_entra_auth_token* patterns with a unified credential parameter and a shared credential→token-provider resolver.
Changes:
- Introduces
AzureCredentialTypes,AzureTokenProvider, andresolve_credential_to_token_provider()for consistent token-provider creation and refresh behavior. - Updates Azure OpenAI clients (chat/responses/assistants) to accept
credential(credential objects or token providers) and always prefer token providers for refresh/caching. - Updates Purview middleware/client and Azure AI / Azure AI Search packages to accept the unified credential type and adjusts tests/exports accordingly.
Reviewed changes
Copilot reviewed 16 out of 16 changed files in this pull request and generated 10 comments.
Show a summary per file
| File | Description |
|---|---|
| python/packages/purview/agent_framework_purview/_middleware.py | Updates Purview middleware to accept unified credential/token-provider types and refreshes docs. |
| python/packages/purview/agent_framework_purview/_client.py | Expands Purview client auth to accept callable token providers in addition to credential objects. |
| python/packages/core/tests/azure/test_entra_id_authentication.py | Replaces tests for token-fetch helpers with tests for credential→token-provider resolution. |
| python/packages/core/tests/azure/test_azure_assistants_client.py | Updates Assistants client tests to validate new credential resolution behavior and new error messaging. |
| python/packages/core/agent_framework/azure/_shared.py | Centralizes OpenAI Azure client construction around resolving credentials into an azure_ad_token_provider. |
| python/packages/core/agent_framework/azure/_responses_client.py | Updates Responses client to accept unified credential/token-provider types and uses them in both direct and project-client paths. |
| python/packages/core/agent_framework/azure/_entra_id_authentication.py | Adds new type aliases and resolve_credential_to_token_provider() using get_bearer_token_provider. |
| python/packages/core/agent_framework/azure/_chat_client.py | Switches Chat client auth inputs to unified credential approach (credential object or token provider). |
| python/packages/core/agent_framework/azure/_assistants_client.py | Switches Assistants client auth inputs to unified credential approach (credential object or token provider). |
| python/packages/core/agent_framework/azure/init.pyi | Updates public typings to export the new credential/token-provider aliases. |
| python/packages/core/agent_framework/azure/init.py | Adds lazy-export entries for the new credential/token-provider aliases; removes old helper export. |
| python/packages/azure-ai/agent_framework_azure_ai/_project_provider.py | Broadens provider’s credential annotation to the shared alias and updates parameter docs. |
| python/packages/azure-ai/agent_framework_azure_ai/_client.py | Broadens client’s credential annotation to the shared alias and updates parameter docs. |
| python/packages/azure-ai/agent_framework_azure_ai/_chat_client.py | Broadens agent chat client’s credential annotation to the shared alias and updates parameter docs. |
| python/packages/azure-ai/agent_framework_azure_ai/_agent_provider.py | Broadens agent provider’s credential annotation to the shared alias and updates parameter docs. |
| python/packages/azure-ai-search/agent_framework_azure_ai_search/_context_provider.py | Broadens Search context provider’s credential annotation to the shared alias and updates parameter docs. |
| @@ -60,10 +62,14 @@ async def close(self) -> None: | |||
| await self._client.aclose() | |||
|
|
|||
| async def _get_token(self, *, tenant_id: str | None = None) -> str: | |||
| """Acquire an access token using either async or sync credential.""" | |||
| scopes = get_purview_scopes(self._settings) | |||
| """Acquire an access token using either async or sync credential, or callable token provider.""" | |||
| cred = self._credential | |||
| token = cred.get_token(*scopes, tenant_id=tenant_id) | |||
| # Callable token provider — returns a token string directly | |||
| if callable(cred) and not isinstance(cred, (TokenCredential, AsyncTokenCredential)): | |||
| result = cred() | |||
| return await result if inspect.isawaitable(result) else result # type: ignore[return-value] | |||
| scopes = get_purview_scopes(self._settings) | |||
| token = cred.get_token(*scopes, tenant_id=tenant_id) # type: ignore[union-attr] | |||
| token = await token if inspect.isawaitable(token) else token | |||
| return token.token | |||
There was a problem hiding this comment.
PurviewClient._get_token() calls sync TokenCredential.get_token() directly on the event loop. This contradicts the class docstring (“invoked in a thread”) and can block the event loop during token acquisition. Consider detecting TokenCredential and running get_token in a worker thread (or switch to always using a bearer token provider) to keep token acquisition non-blocking.
| credential: Azure credential for authentication. Accepts a TokenCredential, | ||
| AsyncTokenCredential, or a callable token provider. |
There was a problem hiding this comment.
The docstring says credential accepts “TokenCredential, AsyncTokenCredential, or a callable token provider”, but the type is AzureCredentialTypes | None (no token provider). Either update the docstring to match (credential objects only) or expand the accepted types and implement a safe adapter for callables.
| credential: Azure credential for authentication. Accepts a TokenCredential, | |
| AsyncTokenCredential, or a callable token provider. | |
| credential: Azure credential for authentication. Accepts an AzureCredentialTypes | |
| instance, such as a TokenCredential or AsyncTokenCredential. |
| resolved_credential: AzureKeyCredential | AsyncTokenCredential | ||
| if credential: | ||
| resolved_credential = credential | ||
| resolved_credential = credential # type: ignore[assignment] |
There was a problem hiding this comment.
AzureAISearchContextProvider now types credential as AzureCredentialTypes (includes sync TokenCredential) and the docstring also mentions callable token providers, but the implementation constructs azure.search.documents.aio.SearchClient and annotates resolved_credential as AzureKeyCredential | AsyncTokenCredential. This is internally inconsistent (note the # type: ignore[assignment]) and may allow passing unsupported credential types to the async SearchClient. Recommend (a) removing the “callable token provider” claim, and (b) narrowing credential back to the credential type(s) that the async SearchClient actually supports, so the type: ignore and mismatched annotations can be removed.
| resolved_credential = credential # type: ignore[assignment] | |
| if isinstance(credential, (AzureKeyCredential, AsyncTokenCredential)): | |
| resolved_credential = credential | |
| else: | |
| raise ServiceInitializationError( | |
| "Unsupported credential type for Azure AI Search async client. " | |
| "Use AzureKeyCredential or an AsyncTokenCredential." | |
| ) |
| credential: Azure credential for authentication. Accepts a TokenCredential, | ||
| AsyncTokenCredential, or a callable token provider. |
There was a problem hiding this comment.
The docstring states credential can be a “callable token provider”, but the parameter type is AzureCredentialTypes | None (credentials only). Please update the docstring to reflect the actual supported types, or expand the accepted types and add explicit support (right now the callable case isn’t handled).
| credential: Azure credential for authentication. Accepts a TokenCredential, | |
| AsyncTokenCredential, or a callable token provider. | |
| credential: Azure credential for authentication. Accepts types compatible with | |
| AzureCredentialTypes (for example, a TokenCredential or AsyncTokenCredential). |
| credential: Azure credential for authentication. Accepts a TokenCredential, | ||
| AsyncTokenCredential, or a callable token provider. |
There was a problem hiding this comment.
The docstring claims credential accepts a “callable token provider”, but the parameter type is AzureCredentialTypes | None. This mismatch is user-facing and can lead to runtime errors if callers pass a callable. Either remove the callable claim or update the signature and implement a supported adapter type.
| credential: Azure credential for authentication. Accepts a TokenCredential, | |
| AsyncTokenCredential, or a callable token provider. | |
| credential: Azure credential for authentication. Accepts a TokenCredential | |
| or AsyncTokenCredential. |
| def _create_client_from_project( | ||
| *, | ||
| project_client: AIProjectClient | None, | ||
| project_endpoint: str | None, | ||
| credential: TokenCredential | None, | ||
| credential: AzureCredentialTypes | AzureTokenProvider | None, | ||
| ) -> AsyncOpenAI: | ||
| """Create an AsyncOpenAI client from an Azure AI Foundry project. | ||
|
|
||
| Args: | ||
| project_client: An existing AIProjectClient to use. | ||
| project_endpoint: The Azure AI Foundry project endpoint URL. | ||
| credential: Azure credential for authentication. | ||
|
|
||
| Returns: | ||
| An AsyncAzureOpenAI client obtained from the project client. | ||
|
|
||
| Raises: | ||
| ServiceInitializationError: If required parameters are missing or | ||
| the azure-ai-projects package is not installed. | ||
| """ | ||
| if project_client is not None: | ||
| return project_client.get_openai_client() | ||
|
|
||
| if not project_endpoint: | ||
| raise ServiceInitializationError( | ||
| "Azure AI project endpoint is required when project_client is not provided." | ||
| ) | ||
| if not credential: | ||
| raise ServiceInitializationError( | ||
| "Azure credential is required when using project_endpoint without a project_client." | ||
| ) | ||
| project_client = AIProjectClient( | ||
| endpoint=project_endpoint, | ||
| credential=credential, # type: ignore[arg-type] | ||
| user_agent=AGENT_FRAMEWORK_USER_AGENT, | ||
| ) |
There was a problem hiding this comment.
AzureOpenAIResponsesClient._create_client_from_project() accepts credential: AzureCredentialTypes | AzureTokenProvider, but then passes it to azure.ai.projects.aio.AIProjectClient with # type: ignore[arg-type]. Token providers (and possibly sync TokenCredential) are unlikely to be valid inputs for the async AIProjectClient, so this path can break at runtime. Recommend validating that credential is an Azure credential object supported by AIProjectClient when using project_endpoint, and raising a clear initialization error otherwise (or narrowing the type for this path).
| # Use provided credential | ||
| if not credential: | ||
| raise ServiceInitializationError("Azure credential is required when project_client is not provided.") | ||
| project_client = AIProjectClient( | ||
| endpoint=resolved_endpoint, | ||
| credential=credential, | ||
| credential=credential, # type: ignore[arg-type] | ||
| user_agent=AGENT_FRAMEWORK_USER_AGENT, | ||
| ) |
There was a problem hiding this comment.
credential is now typed as AzureCredentialTypes (includes sync TokenCredential), but it’s passed into the async AIProjectClient constructor with # type: ignore[arg-type]. This suggests the downstream SDK expects a narrower credential type (likely AsyncTokenCredential). To avoid runtime auth failures and to keep typing honest, either restrict this parameter back to the supported async credential type or add explicit runtime handling/conversion instead of ignoring the type mismatch.
| credential: Azure credential for authentication. Accepts a TokenCredential, | ||
| AsyncTokenCredential, or a callable token provider. |
There was a problem hiding this comment.
The docstring says credential accepts “TokenCredential, AsyncTokenCredential, or a callable token provider”, but the parameter is typed as AzureCredentialTypes | None (no token provider). Please align the docs with the actual supported types to avoid misleading users.
| credential: Azure credential for authentication. Accepts a TokenCredential, | |
| AsyncTokenCredential, or a callable token provider. | |
| credential: Azure credential for authentication. Accepts credentials supported by | |
| AzureCredentialTypes (for example, TokenCredential or AsyncTokenCredential). |
| if not credential: | ||
| raise ServiceInitializationError("Azure credential is required when project_client is not provided.") | ||
|
|
||
| project_client = AIProjectClient( | ||
| endpoint=resolved_endpoint, | ||
| credential=credential, | ||
| credential=credential, # type: ignore[arg-type] | ||
| user_agent=AGENT_FRAMEWORK_USER_AGENT, | ||
| ) |
There was a problem hiding this comment.
credential is now AzureCredentialTypes (includes sync TokenCredential) but is passed into azure.ai.projects.aio.AIProjectClient with # type: ignore[arg-type]. If the SDK only supports AsyncTokenCredential for its async client, accepting TokenCredential here will be a breaking runtime bug. Recommend narrowing the type to the actually supported credential type(s) or adding explicit handling instead of suppressing the type checker.
| if not credential: | ||
| raise ServiceInitializationError("Azure credential is required when agents_client is not provided.") | ||
| self._agents_client = AgentsClient( | ||
| endpoint=resolved_endpoint, | ||
| credential=credential, | ||
| credential=credential, # type: ignore[arg-type] | ||
| user_agent=AGENT_FRAMEWORK_USER_AGENT, | ||
| ) |
There was a problem hiding this comment.
credential is typed as AzureCredentialTypes (includes sync TokenCredential) but is passed to the async AgentsClient constructor with # type: ignore[arg-type]. To avoid accepting unsupported credential types, narrow credential to what AgentsClient actually supports (likely AsyncTokenCredential) or add explicit handling instead of suppressing the type error.
2 participants
Summary
Replace
ad_token,ad_token_provider, andget_entra_auth_tokenwith a unifiedcredentialparameter across all Azure-related Python packages, usingazure.identity.get_bearer_token_providerfor automatic token caching and refresh.Motivation
AzureOpenAIResponsesClientonly accepted syncTokenCredential, failing at runtime in async contexts.ad_token_providerwhen providing Azure credential #3500: Credential handling used one-shot token fetches instead of a token provider with auto-refresh.Changes
New type aliases (
_entra_id_authentication.py)AzureCredentialTypes = TokenCredential | AsyncTokenCredential— credential objectsAzureTokenProvider = Callable[[], str | Awaitable[str]]— callable token providersresolve_credential_to_token_provider()— wraps credentials usingazure.identity.get_bearer_token_providerfor auto-refreshCore package
AzureOpenAIChatClient,AzureOpenAIResponsesClient,AzureOpenAIAssistantsClient: Acceptcredential: AzureCredentialTypes | AzureTokenProviderad_token,ad_token_providerparameters,get_entra_auth_token/get_entra_auth_token_asynchelpersazure-ai package
AzureAIClient,AzureAIAgentClient,AzureAIProjectAgentProvider,AzureAIAgentsProvider: Acceptcredential: AzureCredentialTypes(credential objects only, since the Azure SDK clients handle token refresh internally)azure-ai-search package
AzureAISearchContextProvider: Acceptcredential: AzureCredentialTypespurview package
PurviewClient,PurviewPolicyMiddleware,PurviewChatPolicyMiddleware: Acceptcredential: AzureCredentialTypes | AzureTokenProviderBreaking Changes
ad_tokenandad_token_providerparameters from all Azure OpenAI clientsget_entra_auth_tokenandget_entra_auth_token_asynchelper functionscredential=parameter instead (accepts sync/async credentials and callable token providers)Testing
All existing tests pass across all affected packages (core, azure-ai, azure-ai-search, purview). All fmt, lint, pyright, and mypy checks pass.
Fixes #3449
Fixes #3500