feat(azure-auth): native AAD support for azure_ai/* targets (Foundry hosted agents) - #252
Merged
Merged
Conversation
tangym
requested review from
AaronAspinwall123,
changliu2,
jakepresent and
minthigpen
as code owners
June 24, 2026 21:20
Adds AZURE_FOUNDRY_SCOPE and lets get_azure_token_provider(scope=...) mint tokens for arbitrary audiences. Each scope gets its own cached bearer-token callable; all share a single DefaultAzureCredential instance so adding a second scope does not pay for a second credential-chain probe. Default scope (Azure OpenAI / Cognitive Services) is unchanged, so azure/* callers keep their byte-identical behavior. The existing 'missing dep is cached' contract is preserved across all scopes via the _IDENTITY_IMPORT_ATTEMPTED short-circuit. No call sites updated yet — the family-aware injection that consumes the new scope lands in the next commit.
Adds family-aware AAD injection so azure_ai/agents/<id> (and other azure_ai/* LiteLLM routes) get the same managed-identity / az login / service-principal experience that azure/* models already have. Users no longer need to manually run `az account get-access-token` and paste the result into AZURE_AI_API_KEY before each Foundry-agent eval. The azure_ai branch differs from the azure branch in how the credential reaches LiteLLM: the azure_ai/agents provider only accepts a static api_key string (it sets Authorization: Bearer <api_key> directly), so we call the per-scope token provider once per request and stuff the result into payload['api_key']. The underlying DefaultAzureCredential caches the token, so the per-request call resolves from memory until the token nears expiry. Service Principal env vars (AZURE_TENANT_ID/AZURE_CLIENT_ID/ AZURE_CLIENT_SECRET) flow through the same DefaultAzureCredential chain that handles az login, so production deployments get identical behavior without any ASSERT-side configuration. Behavior for azure/* is byte-identical \u2014 same scope, same payload key (azure_ad_token_provider callable), same precedence (explicit user extra_kwargs still win).
Three improvements that turn the raw LiteLLM error messages into
actionable hints when the failing call targets azure_ai/*:
1. Widen the existing Azure auth-error branch so 401s on azure_ai/*
(not just azure/*) trigger the AAD/install/RBAC hint. The Foundry
variant points at the 'Azure AI User' role and AZURE_AI_API_KEY
(instead of 'Cognitive Services OpenAI User' and AZURE_API_KEY,
which would send users to the wrong resource).
2. Pre-catch the two LiteLLM-side validation errors that the
azure_ai/agents provider raises as APIConnectionError (because
the request never leaves the process):
- 'api_key (Azure AD token) is required ...' -> LLMAuthError
with the install + az login + AZURE_AI_API_KEY hint.
- 'api_base is required for Azure AI Agents' -> LLMInputError
pointing at AZURE_AI_API_BASE.
Both are non-transient and used to be misclassified as
LLMProviderError, which made them eligible for the retry/backoff
loop and buried the real cause behind 5 retry log lines.
Non-azure_ai families are unaffected: the existing tests for
azure/* and non-Azure models still pass.
The 'Azure OpenAI auth mode' / 'Azure OpenAI auth:' log strings predate the Foundry-agent support added in the previous commits and would mislead users who only see azure_ai/* model calls in their config. The mode they describe applies uniformly to both families because the underlying DefaultAzureCredential chain is the same. Rename to plain 'Azure auth mode' / 'Azure auth:'. No test changes needed \u2014 no test pinned the old strings.
Foundry-hosted agents own their tools and instructions server-side, so target.tools and target.system_prompt are silently ignored at runtime for the azure_ai/agents/<id> route. Accept them at parse time and the user gets a successful eval that scores a different agent than they thought they configured. Fail fast at TargetConfig construction with a message that names the right field and the underlying reason. Other azure_ai/* routes (chat completions, embeddings) are runtime-owned just like azure/* and remain unaffected.
Two-line addition to the existing managed-identity section and a single bullet in the schema reference. No new doc page or guide. The Foundry-agent route reuses the same AAD precedence as azure/* models; the only delta is that LiteLLM reads AZURE_AI_API_BASE (Foundry project endpoint) instead of AZURE_API_BASE (Azure OpenAI resource endpoint).
tangym
force-pushed
the
feat/foundry-agent-target
branch
from
June 25, 2026 02:16
1f97de8 to
e0f4eb0
Compare
MohammadHaroonAbuomar
approved these changes
Jun 25, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Adds native AAD support for Azure AI Foundry hosted agents (and other
azure_ai/*LiteLLM routes), so users can pointtarget.modelatazure_ai/agents/<asst_id>and have it just work withaz login/ managed identity / service principal, the same wayazure/*models already do.Problem
target.model: azure_ai/agents/<id>was technically reachable via LiteLLM, but the existing AAD wiring only injected tokens forazure/*models. Users hit two confusing dead ends:AZURE_API_KEY(Azure OpenAI) silently suppressed AAD injection forazure_ai/*calls because mode resolution was global. With a key set for orchestration (azure/gpt-4o-minisystematize/tester/judge), the Foundry call also entered key mode and was sent to Foundry without a bearer token, returning a misleading 'no api_key' error.api_baseerror was misclassified as retryable, so the eval burned 5 retries before surfacing the real cause.Changes
AZURE_FOUNDRY_SCOPE. Per-scopeget_azure_token_provider(scope=...)cache that shares oneDefaultAzureCredential. Family-awareresolve_azure_auth_mode(family=...)that checksAZURE_API_KEYfor theazurefamily andAZURE_AI_API_KEYfor theazure_aifamily. Broaden boot-log label toAzure auth mode:since it covers both families._aad_scope_for_model()dispatch. Inject AAD credentials for both Azure families:azure/*keepsazure_ad_token_providercallable,azure_ai/*calls the provider once per request and stuffs the bearer intopayload['api_key'](what LiteLLM'sazure_ai/agentsprovider expects). Re-resolve mode per family onazure_ai/*injection to avoid theAZURE_API_KEYleak. Classifier widened so 401s and missing-token/AZURE_AI_API_BASEerrors onazure_ai/*surface actionable hints instead of generic retryable errors.target.toolsandtarget.system_promptonazure_ai/agents/*(the hosted agent owns its tools and instructions server-side).azure_ai/agents/<AGENT_ID>.AZURE_API_KEYleak, both LiteLLMapi_baseerror variants, hosted-agent validation guards.Commits (6, each independently green)
924eced54b3f91b0a2699d234429314d7b23b57a1f97de8dUser-visible change
Before this PR, the only paths to a hosted Foundry agent were a callable wrapper or a manually-minted
AZURE_AI_API_KEY=$(az account get-access-token ...). After this PR, pointtarget.modelat the agent and ASSERT mints the AAD token for you:Service Principal (
AZURE_TENANT_ID/AZURE_CLIENT_ID/AZURE_CLIENT_SECRET) and managed identity flow through the sameDefaultAzureCredentialchain.Testing
https://ai.azure.com/.defaultscope, injected intopayload['api_key'], request reached the Foundry endpoint, server-side run completed, reply returned correctly.Out of scope
{project}/agents/{uuid}/threads/...). LiteLLM'sazure_ai/agentsprovider only speaks the v1 Assistants surface (asst_xxx) and rejects UUIDs withInvalid 'assistant_id': '<uuid>'. Expected an ID that begins with 'asst'. v2 agents are evaluable today via the callable-wrapper path demonstrated in Add LangGraph Foundry hosted-agent example with Adaptive Eval #250. A native v2 target will be tracked in a separate issue.Related