Python: forward Azure AI Search query-source identity

[Evan Mattson (moonbox3)]

Motivation & Context

Agentic Azure AI Search retrieval needs the caller's Microsoft Entra identity to enforce
document-level permissions at query time. AzureAISearchContextProvider currently authenticates
the application to the search service but has no supported way to forward a separate user token,
forcing applications to patch the retrieval client themselves.

Description & Review Guide

• What are the major changes?
  • Add an optional async query_source_credential to AzureAISearchContextProvider.
  • Request an Azure AI Search token for each agentic retrieval and forward it through the
    x-ms-query-source-authorization HTTP header.
  • Fail closed with an actionable error when the installed stable Azure Search SDK does not
    support this preview capability.
  • Document the preview requirement and add public-seam tests through before_run(...), including
    a real Azure SDK request-builder and transport-boundary assertion.
• What is the impact of these changes?
  • Applications using permission-aware knowledge sources can opt in to caller-specific result
    trimming without changing existing API-key or managed-identity behavior.
  • Existing callers are unaffected because the new credential is optional.
• What do you want reviewers to focus on?
  • The preview capability gate and fail-closed behavior.
  • The separation between the service credential and the per-query caller identity header.

Related Issue

Fixes #6864

This supersedes #6915 by carrying the feature onto current main and adding stable-SDK
fail-closed behavior plus transport-level regression coverage.

Contribution Checklist

• The code builds clean without any errors or warnings
• All unit tests pass, and I have added new tests where possible
• The PR follows the Contribution Guidelines
• This PR is linked to an issue and there is no other open PR for this issue (see Related Issue above).
• This is not a breaking change. If it is a breaking change, add the breaking change label (or add "[BREAKING]" to the title prefix, before or after any language prefix) — a workflow keeps the label and title prefix in sync automatically.

[Copilot]

Pull request overview

This PR adds an opt-in mechanism for forwarding a caller’s Microsoft Entra identity to Azure AI Search during agentic Knowledge Base retrieval, enabling document-level permission trimming at query time while keeping existing API key / managed identity behavior unchanged by default.

Changes:

• Adds an optional async query_source_credential parameter to AzureAISearchContextProvider and uses it to fetch a Search-scoped token per agentic retrieval.
• Forwards the token via the x-ms-query-source-authorization HTTP header, and fails closed with an actionable error when the installed SDK channel doesn’t support the capability.
• Documents the preview requirement and adds transport-boundary tests asserting the header is emitted (and that stable/unsupported scenarios fail before any HTTP transport is reached).

Reviewed changes

Copilot reviewed 4 out of 4 changed files in this pull request and generated no comments.

File	Description
python/packages/azure-ai-search/agent_framework_azure_ai_search/_context_provider.py	Introduces query_source_credential and injects x-ms-query-source-authorization on agentic retrieval; adds preview gating + fail-closed behavior.
python/packages/azure-ai-search/tests/test_aisearch_context_provider.py	Adds before_run-based tests covering preview gating, header emission at the transport boundary, omission when unset, and rejection of sync credentials.
python/packages/azure-ai-search/README.md	Documents the new query-time identity option and its preview SDK requirement with a usage snippet.
python/packages/azure-ai-search/AGENTS.md	Notes the preview-only gate and how query_source_credential is used to forward the per-request identity header.

[github-actions]

Automated Code Review

Reviewers: 5 | Confidence: 90% | Result: All clear

Reviewed: Correctness, Security Reliability, Test Coverage, Failure Modes, Design Approach

---

Automated review by moonbox3's agents

[github-actions]

Python Test Coverage Report •

File	Stmts	Miss	Cover	Missing
packages/azure-ai-search/agent_framework_azure_ai_search
_context_provider.py	409	17	95%	99–100, 125–128, 560, 639–640, 768–769, 867, 872, 877, 989–990, 1038
TOTAL	47360	4527	90%

Python Unit Test Overview

Tests	Skipped	Failures	Errors	Time
9748	34 💤	0 ❌	0 🔥	2m 31s ⏱️

[github-actions]

Automated Code Review

Reviewers: 5 | Confidence: 90%

✓ Correctness

The PR cleanly adds an optional query_source_credential to AzureAISearchContextProvider forwarding caller-specific Azure AI Search identity. The fail-closed gate (_query_source_authorization_available) correctly prevents use on the stable SDK, the inspect.isawaitable check detects accidentally-passed sync credentials at runtime, and the headers kwarg follows standard Azure SDK pipeline patterns. All constructor overloads consistently include the new parameter with correct types and defaults. No correctness issues found.

✓ Security Reliability

This PR adds optional query-time user identity forwarding for agentic Azure AI Search retrieval. The security design is sound: it fails closed when the preview SDK is absent (checked before any HTTP transport), validates async credentials at runtime via inspect.isawaitable, uses the correct Azure Search resource scope, and does not log token values. The token is only forwarded after both the SDK capability gate and the async credential check pass. No significant security or reliability issues found.

✓ Test Coverage

Test coverage for the new query_source_credential feature is solid, covering the four critical paths: preview-SDK fail-closed gate, HTTP header serialization via real SDK transport, header omission without credential, and sync credential rejection. One medium-priority gap: the happy-path async test verifies the header value but doesn't assert the resource scope passed to get_token, relying on the sync rejection test for scope verification. A minor gap exists around semantic-mode behavior with query_source_credential set (silently ignored, documented but untested).

✓ Failure Modes

The PR introduces query-source identity forwarding for agentic Azure AI Search retrieval. The implementation correctly fails closed when the preview SDK is missing (ValueError before any I/O), detects sync credentials at runtime with a TypeError before retrieval, and propagates token acquisition failures loudly. No silent failure paths, swallowed exceptions, or operational failure modes were identified. The gating, error handling, and test coverage are well-structured.

✗ Design Approach

The main design issue is the placement of the new query_source_credential validation. In the index-backed agentic path, the provider now provisions or updates Azure AI Search resources before it proves that the caller-suplied query credential is usable, so invalid input can still mutate remote state before the request fails.

Flagged Issues

• In _context_provider.py, _ensure_knowledge_base() (line 894) runs before query_source_credential validation (lines 939-943). In the index-backed agentic flow, _ensure_knowledge_base() can create a knowledge source and create/update the knowledge base (lines 821-874), so passing an invalid sync credential mutates server state before the TypeError is raised. Move credential validation before _ensure_knowledge_base() so invalid caller identity never triggers provisioning side effects.

---

Automated review by moonbox3's agents

[github-actions]

Flagged issue

In _context_provider.py, _ensure_knowledge_base() (line 894) runs before query_source_credential validation (lines 939-943). In the index-backed agentic flow, _ensure_knowledge_base() can create a knowledge source and create/update the knowledge base (lines 821-874), so passing an invalid sync credential mutates server state before the TypeError is raised. Move credential validation before _ensure_knowledge_base() so invalid caller identity never triggers provisioning side effects.

---

Source: automated DevFlow PR review

[github-actions]

Automated Code Review

Reviewers: 5 | Confidence: 70% | Result: All clear

Reviewed: Correctness, Security Reliability, Test Coverage, Failure Modes, Design Approach

---

Automated review by moonbox3's agents