Python: forward Azure AI Search query-source identity

[Karunanidhi Mishra (kmishra1204)]

Summary

• add an optional query_source_credential to AzureAISearchContextProvider
• forward the caller's Azure AI Search authorization token on agentic Knowledge Base retrieval requests
• document the query-time identity option, add regression coverage, and record the change in the Python changelog

This addresses the Python side of #6864. The option is explicit so existing API key and managed identity flows keep their current behavior unless a caller supplies an async query-source credential.

Validation

• uv run --project python --dev --no-sync pytest -q python/packages/azure-ai-search/tests/test_aisearch_context_provider.py::TestAgenticSearch::test_minimal_reasoning_returns_results python/packages/azure-ai-search/tests/test_aisearch_context_provider.py::TestAgenticSearch::test_query_source_credential_forwards_authorization_token python/packages/azure-ai-search/tests/test_aisearch_context_provider.py::TestAgenticSearch::test_query_source_credential_requires_async_credential
• uv run --project python --dev --no-sync ruff check python/packages/azure-ai-search/agent_framework_azure_ai_search/_context_provider.py python/packages/azure-ai-search/tests/test_aisearch_context_provider.py
• uv run --project python --dev --no-sync pyright python/packages/azure-ai-search/agent_framework_azure_ai_search/_context_provider.py
• git diff --check -- python/CHANGELOG.md python/packages/azure-ai-search/README.md python/packages/azure-ai-search/agent_framework_azure_ai_search/_context_provider.py python/packages/azure-ai-search/tests/test_aisearch_context_provider.py

[Copilot]

Pull request overview

Adds an opt-in way for Python agentic Azure AI Search retrieval to forward a caller’s query-time identity (Entra token) to Knowledge Base retrieval requests, enabling document-level permission trimming scenarios without changing existing API key/managed identity defaults.

Changes:

• Added an optional query_source_credential parameter to AzureAISearchContextProvider and forwarded its token to KnowledgeBaseRetrievalClient.retrieve(...) via x_ms_query_source_authorization.
• Documented query-time identity forwarding in the package README.
• Added regression tests covering both default behavior (no header) and forwarding behavior (header present + correct scope).

Reviewed changes

Copilot reviewed 3 out of 3 changed files in this pull request and generated 2 comments.

File	Description
python/packages/azure-ai-search/agent_framework_azure_ai_search/_context_provider.py	Adds query_source_credential, obtains the Search scope token, and forwards it on agentic retrieval calls.
python/packages/azure-ai-search/tests/test_aisearch_context_provider.py	Adds assertions and a new test verifying the authorization header forwarding behavior.
python/packages/azure-ai-search/README.md	Documents the query-time user identity option and shows how to pass query_source_credential.

[Karunanidhi Mishra (kmishra1204)]

Updated in bd48f59 to address the async credential feedback: the docs now say query_source_credential must be an async Azure credential, _query_source_authorization() raises a clear TypeError if a sync credential is passed, and there is regression coverage for that misuse case.

[github-actions]

Python Test Coverage Report •

File	Stmts	Miss	Cover	Missing
packages/azure-ai-search/agent_framework_azure_ai_search
_context_provider.py	396	17	95%	95–96, 121–124, 547, 625–626, 753–754, 844, 849, 854, 953–954, 1002
TOTAL	44017	5303	87%

Python Unit Test Overview

Tests	Skipped	Failures	Errors	Time
8595	33 💤	0 ❌	0 🔥	2m 18s ⏱️

[Evan Mattson (moonbox3)]

Please also have a look at the failing CI/CD (code quality) checks.

[Karunanidhi Mishra (kmishra1204)]

Thanks for the CI note. I rebased this branch on current main and addressed the test typing failures from the Test Typing Checks (3.11) job.

What changed:

• Stored mock_retrieval.retrieve.await_args in a local variable and asserted it is not None before reading .kwargs in the new Azure AI Search tests.
• The rebase also picks up upstream #6944, which fixes the root README stargazers link that was failing markdown-link-check.

Validation run locally:

• uv run poe test-typing -P azure-ai-search --checker ty --checker zuban -> passed
• uv run --directory python pytest -m "not integration" packages\azure-ai-search\tests\test_aisearch_context_provider.py -k "query_source_credential or agentic_search" -> 3 passed
• uv run poe lint -P azure-ai-search -> passed
• uv run poe pyright -P azure-ai-search -> passed
• git diff --check -> passed

I also ran uv run poe test-typing -P azure-ai-search; the locally installed environment stops in mypy on a NumPy stub requiring Python 3.12+, before checking this package. The CI mypy job was already green, and the CI failures called out in this PR were the ty/zuban diagnostics fixed above.

[Karunanidhi Mishra (kmishra1204)]

Rebased this branch onto current upstream main to clear the behind-base state.

Current head: 38b5ac45

Validation rerun locally:

• uv run poe test-typing -P azure-ai-search --checker ty --checker zuban
• uv run pytest packages/azure-ai-search/tests/test_aisearch_context_provider.py -k "query_source or credential or retrieve"
• uv run poe lint -P azure-ai-search
• git diff --check upstream/main...HEAD

Results:

• Type checks passed (ty and zuban).
• Targeted tests passed: 5 passed, 123 deselected.
• Package lint passed.
• Diff check passed.

Live status after push: code checks have no failures. The remaining visible gates are the merge gatekeeper still in progress and the queued license/cla policy check.

[Karunanidhi Mishra (kmishra1204)]

Updated in 1358f3bc for the review note on the retrieval kwarg.

Change made:

• retrieve(...) now passes query_source_authorization=... instead of the generic x_ms_query_source_authorization kwarg.
• The focused regression assertions now check the SDK-supported kwarg name.

Validation rerun:

• uv run pytest packages/azure-ai-search/tests/test_aisearch_context_provider.py -k "query_source or credential or retrieve" -> 5 passed, 123 deselected
• uv run poe lint -P azure-ai-search -> passed
• uv run poe test-typing -P azure-ai-search --checker ty --checker zuban -> passed
• git diff --check upstream/main...HEAD -> passed

[Karunanidhi Mishra (kmishra1204)]

Also removed the manual python/CHANGELOG.md entry in 97f8a36a, per the release-note guidance.

Current review-response commits:

• 1358f3bc uses query_source_authorization instead of the generic x_ms_* kwarg.
• 97f8a36a removes the manual changelog entry.

Validation after the latest push:

• uv run pytest packages/azure-ai-search/tests/test_aisearch_context_provider.py -k "query_source or credential or retrieve" -> 5 passed, 123 deselected
• git diff --check upstream/main...HEAD -> passed

I resolved the two review threads that these commits address.

[Evan Mattson (moonbox3)]

Karunanidhi Mishra (@kmishra1204) please accept the CLA

[Evan Mattson (moonbox3)]

Karunanidhi Mishra (@kmishra1204) ping to accept the CLA please

[Evan Mattson (moonbox3)]

We appreciate your work on this. We need to get this fix in, so superseded by #7278.