From e460b2f7d6c27d365e21333757880ce535542d2c Mon Sep 17 00:00:00 2001 From: Karunanidhi Mishra <11963379+kmishra1204@users.noreply.github.com> Date: Sat, 4 Jul 2026 17:31:05 -0500 Subject: [PATCH 1/6] Python: forward Azure AI Search query-source identity --- python/packages/azure-ai-search/README.md | 17 ++++++++++ .../_context_provider.py | 24 +++++++++++++- .../tests/test_aisearch_context_provider.py | 33 +++++++++++++++++++ 3 files changed, 73 insertions(+), 1 deletion(-) diff --git a/python/packages/azure-ai-search/README.md b/python/packages/azure-ai-search/README.md index 234ccdf7f7f..7c0567019b3 100644 --- a/python/packages/azure-ai-search/README.md +++ b/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: diff --git a/python/packages/azure-ai-search/agent_framework_azure_ai_search/_context_provider.py b/python/packages/azure-ai-search/agent_framework_azure_ai_search/_context_provider.py index 347e03c8186..f567d2cc1a8 100644 --- a/python/packages/azure-ai-search/agent_framework_azure_ai_search/_context_provider.py +++ b/python/packages/azure-ai-search/agent_framework_azure_ai_search/_context_provider.py @@ -130,6 +130,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 +197,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 +223,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 +252,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 +278,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 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 +307,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 +333,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 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 +362,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 +392,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 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 +420,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 +449,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 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 +525,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 +898,20 @@ 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["x_ms_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 = await self.query_source_credential.get_token(_AZURE_SEARCH_RESOURCE_SCOPE) + return access_token.token + @staticmethod def _prepare_messages_for_kb_search(messages: list[Message]) -> list[KnowledgeBaseMessage]: """Convert framework Messages to KnowledgeBaseMessages for agentic retrieval. diff --git a/python/packages/azure-ai-search/tests/test_aisearch_context_provider.py b/python/packages/azure-ai-search/tests/test_aisearch_context_provider.py index 17ae6eb06af..90c2ed603fd 100644 --- a/python/packages/azure-ai-search/tests/test_aisearch_context_provider.py +++ b/python/packages/azure-ai-search/tests/test_aisearch_context_provider.py @@ -1351,6 +1351,39 @@ async def test_minimal_reasoning_returns_results(self) -> None: assert len(results) == 1 assert results[0].text == "Answer text" assert results[0].role == "assistant" + assert "x_ms_query_source_authorization" not in mock_retrieval.retrieve.await_args.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") + assert mock_retrieval.retrieve.await_args.kwargs["x_ms_query_source_authorization"] == "user-token" async def test_non_minimal_reasoning_uses_messages(self) -> None: provider = _make_provider() From 69e749c02f3962711d72d3140c67c06767e59f69 Mon Sep 17 00:00:00 2001 From: Karunanidhi Mishra <11963379+kmishra1204@users.noreply.github.com> Date: Sat, 4 Jul 2026 17:40:19 -0500 Subject: [PATCH 2/6] Clarify query source credential must be async --- .../_context_provider.py | 17 ++++++++++++----- .../tests/test_aisearch_context_provider.py | 17 +++++++++++++++++ 2 files changed, 29 insertions(+), 5 deletions(-) diff --git a/python/packages/azure-ai-search/agent_framework_azure_ai_search/_context_provider.py b/python/packages/azure-ai-search/agent_framework_azure_ai_search/_context_provider.py index f567d2cc1a8..7636af9686a 100644 --- a/python/packages/azure-ai-search/agent_framework_azure_ai_search/_context_provider.py +++ b/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 @@ -278,7 +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 Azure credential for per-query user identity forwarding. + 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. @@ -333,7 +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 Azure credential for per-query user identity forwarding. + 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. @@ -392,7 +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 Azure credential for per-query user identity forwarding. + 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. @@ -449,7 +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 Azure credential for per-query user identity forwarding. + 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. @@ -909,7 +910,13 @@ 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 = await self.query_source_credential.get_token(_AZURE_SEARCH_RESOURCE_SCOPE) + 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 diff --git a/python/packages/azure-ai-search/tests/test_aisearch_context_provider.py b/python/packages/azure-ai-search/tests/test_aisearch_context_provider.py index 90c2ed603fd..5b80dd9d1ca 100644 --- a/python/packages/azure-ai-search/tests/test_aisearch_context_provider.py +++ b/python/packages/azure-ai-search/tests/test_aisearch_context_provider.py @@ -1385,6 +1385,23 @@ async def test_query_source_credential_forwards_authorization_token(self) -> Non query_source_credential.get_token.assert_awaited_once_with("https://search.azure.com/.default") assert mock_retrieval.retrieve.await_args.kwargs["x_ms_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() provider._knowledge_base_initialized = True From dc78019e8537d7417acc2f2a744cb1c462705a26 Mon Sep 17 00:00:00 2001 From: Karunanidhi Mishra <11963379+kmishra1204@users.noreply.github.com> Date: Sat, 4 Jul 2026 17:54:08 -0500 Subject: [PATCH 3/6] Add Azure AI Search query source changelog entry --- python/CHANGELOG.md | 3 +++ 1 file changed, 3 insertions(+) diff --git a/python/CHANGELOG.md b/python/CHANGELOG.md index 338e2678b56..f579aa2dd60 100644 --- a/python/CHANGELOG.md +++ b/python/CHANGELOG.md @@ -7,6 +7,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] +### Added +- **agent-framework-azure-ai-search**: Add opt-in query source credential forwarding for agentic Azure AI Search retrieval ([#6915](https://github.com/microsoft/agent-framework/pull/6915)) + ## [1.10.0] - 2026-06-30 ### Added From 38b5ac45da15833042479eddc47fe421a5b88cc5 Mon Sep 17 00:00:00 2001 From: Karunanidhi Mishra <11963379+kmishra1204@users.noreply.github.com> Date: Tue, 7 Jul 2026 00:40:54 -0500 Subject: [PATCH 4/6] Fix Azure AI Search test typing --- .../tests/test_aisearch_context_provider.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/python/packages/azure-ai-search/tests/test_aisearch_context_provider.py b/python/packages/azure-ai-search/tests/test_aisearch_context_provider.py index 5b80dd9d1ca..4e1d88f8aca 100644 --- a/python/packages/azure-ai-search/tests/test_aisearch_context_provider.py +++ b/python/packages/azure-ai-search/tests/test_aisearch_context_provider.py @@ -1351,7 +1351,9 @@ async def test_minimal_reasoning_returns_results(self) -> None: assert len(results) == 1 assert results[0].text == "Answer text" assert results[0].role == "assistant" - assert "x_ms_query_source_authorization" not in mock_retrieval.retrieve.await_args.kwargs + retrieve_call = mock_retrieval.retrieve.await_args + assert retrieve_call is not None + assert "x_ms_query_source_authorization" not in retrieve_call.kwargs async def test_query_source_credential_forwards_authorization_token(self) -> None: query_source_credential = AsyncMock() @@ -1383,7 +1385,9 @@ async def test_query_source_credential_forwards_authorization_token(self) -> Non assert len(results) == 1 assert results[0].text == "Answer text" query_source_credential.get_token.assert_awaited_once_with("https://search.azure.com/.default") - assert mock_retrieval.retrieve.await_args.kwargs["x_ms_query_source_authorization"] == "user-token" + retrieve_call = mock_retrieval.retrieve.await_args + assert retrieve_call is not None + assert retrieve_call.kwargs["x_ms_query_source_authorization"] == "user-token" async def test_query_source_credential_requires_async_credential(self) -> None: query_source_credential = Mock() From 1358f3bcf1ae474c301be21a7aa0ae25bbd09dc8 Mon Sep 17 00:00:00 2001 From: Karunanidhi Mishra <11963379+kmishra1204@users.noreply.github.com> Date: Tue, 7 Jul 2026 03:19:46 -0500 Subject: [PATCH 5/6] fix: use query source authorization parameter --- .../agent_framework_azure_ai_search/_context_provider.py | 2 +- .../azure-ai-search/tests/test_aisearch_context_provider.py | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/python/packages/azure-ai-search/agent_framework_azure_ai_search/_context_provider.py b/python/packages/azure-ai-search/agent_framework_azure_ai_search/_context_provider.py index 7636af9686a..e274523002e 100644 --- a/python/packages/azure-ai-search/agent_framework_azure_ai_search/_context_provider.py +++ b/python/packages/azure-ai-search/agent_framework_azure_ai_search/_context_provider.py @@ -901,7 +901,7 @@ async def _agentic_search(self, messages: list[Message]) -> list[Message]: raise RuntimeError("Retrieval client not initialized.") retrieve_kwargs: dict[str, Any] = {"retrieval_request": retrieval_request} if query_source_authorization := await self._query_source_authorization(): - retrieve_kwargs["x_ms_query_source_authorization"] = 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) diff --git a/python/packages/azure-ai-search/tests/test_aisearch_context_provider.py b/python/packages/azure-ai-search/tests/test_aisearch_context_provider.py index 4e1d88f8aca..4b22b6401db 100644 --- a/python/packages/azure-ai-search/tests/test_aisearch_context_provider.py +++ b/python/packages/azure-ai-search/tests/test_aisearch_context_provider.py @@ -1353,7 +1353,7 @@ async def test_minimal_reasoning_returns_results(self) -> None: assert results[0].role == "assistant" retrieve_call = mock_retrieval.retrieve.await_args assert retrieve_call is not None - assert "x_ms_query_source_authorization" not in retrieve_call.kwargs + assert "query_source_authorization" not in retrieve_call.kwargs async def test_query_source_credential_forwards_authorization_token(self) -> None: query_source_credential = AsyncMock() @@ -1387,7 +1387,7 @@ async def test_query_source_credential_forwards_authorization_token(self) -> Non 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["x_ms_query_source_authorization"] == "user-token" + assert retrieve_call.kwargs["query_source_authorization"] == "user-token" async def test_query_source_credential_requires_async_credential(self) -> None: query_source_credential = Mock() From 97f8a36a2c49d4b8d12b9f53bdb8feec0a847ffb Mon Sep 17 00:00:00 2001 From: Karunanidhi Mishra <11963379+kmishra1204@users.noreply.github.com> Date: Tue, 7 Jul 2026 03:20:38 -0500 Subject: [PATCH 6/6] chore: remove manual azure search changelog entry --- python/CHANGELOG.md | 3 --- 1 file changed, 3 deletions(-) diff --git a/python/CHANGELOG.md b/python/CHANGELOG.md index f579aa2dd60..338e2678b56 100644 --- a/python/CHANGELOG.md +++ b/python/CHANGELOG.md @@ -7,9 +7,6 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] -### Added -- **agent-framework-azure-ai-search**: Add opt-in query source credential forwarding for agentic Azure AI Search retrieval ([#6915](https://github.com/microsoft/agent-framework/pull/6915)) - ## [1.10.0] - 2026-06-30 ### Added