Python: forward Azure AI Search query-source identity

python/packages/azure-ai-search/README.md

@@ -31,6 +31,23 @@ ship only in the preview build. When a stable build is installed, the provider u
 output with minimal reasoning effort and raises an actionable error if a preview-only option is
 explicitly requested. Switching channels is a single change — the install — with no code edits.
 
+### Query-time user identity
+
+Agentic retrieval can forward a caller-specific Azure AI Search authorization token when the
+index uses permission fields for document-level access control. Pass an async credential for the
+caller via `query_source_credential`; the provider requests the Azure AI Search resource scope and
+forwards the token on each Knowledge Base retrieval request.
+
+```python
+context_provider = AzureAISearchContextProvider(
+    endpoint=search_endpoint,
+    credential=application_credential,
+    mode="agentic",
+    knowledge_base_name=knowledge_base_name,
+    query_source_credential=user_credential,
+)
+```
+
 ### Basic Usage Example
 
 See the [Azure AI Search context provider examples](../../samples/02-agents/context_providers/azure_ai_search/) which demonstrate:

python/packages/azure-ai-search/agent_framework_azure_ai_search/_context_provider.py

@@ -9,6 +9,7 @@
 from __future__ import annotations
 
 import importlib.metadata
+import inspect
 import logging
 import sys
 from collections.abc import Awaitable, Callable
@@ -130,6 +131,7 @@
 logger = logging.getLogger("agent_framework.azure_ai_search")
 
 _DEFAULT_AGENTIC_MESSAGE_HISTORY_COUNT = 10
+_AZURE_SEARCH_RESOURCE_SCOPE = "https://search.azure.com/.default"
 
 
 def _installed_search_documents_version() -> str:
@@ -196,6 +198,7 @@ def __init__(
         azure_openai_api_key: str | None = None,
         knowledge_base_output_mode: KnowledgeBaseOutputModeLiteral = "extractive_data",
         retrieval_reasoning_effort: RetrievalReasoningEffortLiteral = "minimal",
+        query_source_credential: AsyncTokenCredential | None = None,
         agentic_message_history_count: int = _DEFAULT_AGENTIC_MESSAGE_HISTORY_COUNT,
         env_file_path: str | None = None,
         env_file_encoding: str | None = None,
@@ -221,6 +224,7 @@ def __init__(
             azure_openai_api_key: Unused in semantic mode.
             knowledge_base_output_mode: Unused in semantic mode.
             retrieval_reasoning_effort: Unused in semantic mode.
+            query_source_credential: Unused in semantic mode.
             agentic_message_history_count: Unused in semantic mode.
             env_file_path: Optional ``.env`` file checked before process environment variables.
             env_file_encoding: Encoding for the ``.env`` file.
@@ -249,6 +253,7 @@ def __init__(
         azure_openai_api_key: str | None = None,
         knowledge_base_output_mode: KnowledgeBaseOutputModeLiteral = "extractive_data",
         retrieval_reasoning_effort: RetrievalReasoningEffortLiteral = "minimal",
+        query_source_credential: AsyncTokenCredential | None = None,
         agentic_message_history_count: int = _DEFAULT_AGENTIC_MESSAGE_HISTORY_COUNT,
         env_file_path: str | None = None,
         env_file_encoding: str | None = None,
@@ -274,6 +279,7 @@ def __init__(
             azure_openai_api_key: Optional Azure OpenAI API key for Knowledge Base creation.
             knowledge_base_output_mode: Output mode for Knowledge Base retrieval.
             retrieval_reasoning_effort: Reasoning effort for query planning.
+            query_source_credential: Optional async Azure credential for per-query user identity forwarding.
             agentic_message_history_count: Number of recent messages included in retrieval.
             env_file_path: Optional ``.env`` file checked before process environment variables.
             env_file_encoding: Encoding for the ``.env`` file.
@@ -302,6 +308,7 @@ def __init__(
         azure_openai_api_key: str | None = None,
         knowledge_base_output_mode: KnowledgeBaseOutputModeLiteral = "extractive_data",
         retrieval_reasoning_effort: RetrievalReasoningEffortLiteral = "minimal",
+        query_source_credential: AsyncTokenCredential | None = None,
         agentic_message_history_count: int = _DEFAULT_AGENTIC_MESSAGE_HISTORY_COUNT,
         env_file_path: str | None = None,
         env_file_encoding: str | None = None,
@@ -327,6 +334,7 @@ def __init__(
             azure_openai_api_key: Unused when connecting to an existing Knowledge Base.
             knowledge_base_output_mode: Output mode for Knowledge Base retrieval.
             retrieval_reasoning_effort: Reasoning effort for query planning.
+            query_source_credential: Optional async Azure credential for per-query user identity forwarding.
             agentic_message_history_count: Number of recent messages included in retrieval.
             env_file_path: Optional ``.env`` file checked before process environment variables.
             env_file_encoding: Encoding for the ``.env`` file.
@@ -355,6 +363,7 @@ def __init__(
         azure_openai_api_key: str | None = None,
         knowledge_base_output_mode: KnowledgeBaseOutputModeLiteral = "extractive_data",
         retrieval_reasoning_effort: RetrievalReasoningEffortLiteral = "minimal",
+        query_source_credential: AsyncTokenCredential | None = None,
         agentic_message_history_count: int = _DEFAULT_AGENTIC_MESSAGE_HISTORY_COUNT,
         env_file_path: str | None = None,
         env_file_encoding: str | None = None,
@@ -384,6 +393,7 @@ def __init__(
             azure_openai_api_key: Optional Azure OpenAI API key for Knowledge Base creation.
             knowledge_base_output_mode: Output mode for Knowledge Base retrieval.
             retrieval_reasoning_effort: Reasoning effort for query planning.
+            query_source_credential: Optional async Azure credential for per-query user identity forwarding.
             agentic_message_history_count: Number of recent messages included in retrieval.
             env_file_path: Optional ``.env`` file checked before process environment variables.
             env_file_encoding: Encoding for the ``.env`` file.
@@ -411,6 +421,7 @@ def __init__(
         azure_openai_api_key: str | None = None,
         knowledge_base_output_mode: KnowledgeBaseOutputModeLiteral = "extractive_data",
         retrieval_reasoning_effort: RetrievalReasoningEffortLiteral = "minimal",
+        query_source_credential: AsyncTokenCredential | None = None,
         agentic_message_history_count: int = _DEFAULT_AGENTIC_MESSAGE_HISTORY_COUNT,
         env_file_path: str | None = None,
         env_file_encoding: str | None = None,
@@ -439,6 +450,7 @@ def __init__(
             azure_openai_api_key: Azure OpenAI API key.
             knowledge_base_output_mode: Output mode for Knowledge Base retrieval.
             retrieval_reasoning_effort: Reasoning effort for Knowledge Base query planning.
+            query_source_credential: Optional async Azure credential for per-query user identity forwarding.
             agentic_message_history_count: Number of recent messages for agentic mode.
             env_file_path: Path to environment file for loading settings.
             env_file_encoding: Encoding of the environment file.
@@ -514,6 +526,7 @@ def __init__(
         self.azure_openai_api_key = azure_openai_api_key
         self.knowledge_base_output_mode = knowledge_base_output_mode
         self.retrieval_reasoning_effort = retrieval_reasoning_effort
+        self.query_source_credential = query_source_credential
         self.agentic_message_history_count = agentic_message_history_count
 
         self._use_existing_knowledge_base = False
@@ -886,10 +899,26 @@ async def _agentic_search(self, messages: list[Message]) -> list[Message]:
 
         if not self._retrieval_client:
             raise RuntimeError("Retrieval client not initialized.")
-        retrieval_result = await self._retrieval_client.retrieve(retrieval_request=retrieval_request)
+        retrieve_kwargs: dict[str, Any] = {"retrieval_request": retrieval_request}
+        if query_source_authorization := await self._query_source_authorization():
+            retrieve_kwargs["query_source_authorization"] = query_source_authorization
+        retrieval_result = await self._retrieval_client.retrieve(**retrieve_kwargs)
 
         return self._parse_messages_from_kb_response(retrieval_result)
 
+    async def _query_source_authorization(self) -> str | None:
+        """Return a per-query Azure AI Search authorization token, when configured."""
+        if self.query_source_credential is None:
+            return None
+        access_token_result = self.query_source_credential.get_token(_AZURE_SEARCH_RESOURCE_SCOPE)
+        if not inspect.isawaitable(access_token_result):
+            raise TypeError(
+                "query_source_credential must be an async Azure credential. "
+                "Pass an azure.core.credentials_async.AsyncTokenCredential."
+            )
+        access_token = await access_token_result
+        return access_token.token
+
     @staticmethod
     def _prepare_messages_for_kb_search(messages: list[Message]) -> list[KnowledgeBaseMessage]:
         """Convert framework Messages to KnowledgeBaseMessages for agentic retrieval.

python/packages/azure-ai-search/tests/test_aisearch_context_provider.py

@@ -1351,6 +1351,60 @@ async def test_minimal_reasoning_returns_results(self) -> None:
         assert len(results) == 1
         assert results[0].text == "Answer text"
         assert results[0].role == "assistant"
+        retrieve_call = mock_retrieval.retrieve.await_args
+        assert retrieve_call is not None
+        assert "query_source_authorization" not in retrieve_call.kwargs
+
+    async def test_query_source_credential_forwards_authorization_token(self) -> None:
+        query_source_credential = AsyncMock()
+        query_source_credential.get_token = AsyncMock(return_value=SimpleNamespace(token="user-token"))
+        provider = _make_provider(query_source_credential=query_source_credential)
+        provider._knowledge_base_initialized = True
+        provider.knowledge_base_name = "kb"
+        provider.retrieval_reasoning_effort = "minimal"
+
+        mock_content = Mock()
+        mock_content.text = "Answer text"
+        mock_message = Mock()
+        mock_message.role = "assistant"
+        mock_message.content = [mock_content]
+        mock_result = Mock()
+        mock_result.response = [mock_message]
+        mock_result.references = None
+
+        mock_retrieval = AsyncMock()
+        mock_retrieval.retrieve = AsyncMock(return_value=mock_result)
+        provider._retrieval_client = mock_retrieval
+
+        with patch(
+            "agent_framework_azure_ai_search._context_provider.KnowledgeBaseMessageTextContent",
+            type(mock_content),
+        ):
+            results = await provider._agentic_search([Message(role="user", contents=["test query"])])
+
+        assert len(results) == 1
+        assert results[0].text == "Answer text"
+        query_source_credential.get_token.assert_awaited_once_with("https://search.azure.com/.default")
+        retrieve_call = mock_retrieval.retrieve.await_args
+        assert retrieve_call is not None
+        assert retrieve_call.kwargs["query_source_authorization"] == "user-token"
+
+    async def test_query_source_credential_requires_async_credential(self) -> None:
+        query_source_credential = Mock()
+        query_source_credential.get_token = Mock(return_value=SimpleNamespace(token="user-token"))
+        provider = _make_provider(query_source_credential=query_source_credential)
+        provider._knowledge_base_initialized = True
+        provider.knowledge_base_name = "kb"
+        provider.retrieval_reasoning_effort = "minimal"
+
+        mock_retrieval = AsyncMock()
+        provider._retrieval_client = mock_retrieval
+
+        with pytest.raises(TypeError, match="query_source_credential must be an async Azure credential"):
+            await provider._agentic_search([Message(role="user", contents=["test query"])])
+
+        query_source_credential.get_token.assert_called_once_with("https://search.azure.com/.default")
+        mock_retrieval.retrieve.assert_not_awaited()
 
     async def test_non_minimal_reasoning_uses_messages(self) -> None:
         provider = _make_provider()