Skip to content

feat(azure-auth): native AAD support for azure_ai/* targets (Foundry hosted agents) - #252

Merged
MohammadHaroonAbuomar merged 6 commits into
mainfrom
feat/foundry-agent-target
Jun 25, 2026
Merged

feat(azure-auth): native AAD support for azure_ai/* targets (Foundry hosted agents)#252
MohammadHaroonAbuomar merged 6 commits into
mainfrom
feat/foundry-agent-target

Conversation

@tangym

@tangym tangym commented Jun 24, 2026

Copy link
Copy Markdown
Collaborator

Summary

Adds native AAD support for Azure AI Foundry hosted agents (and other azure_ai/* LiteLLM routes), so users can point target.model at azure_ai/agents/<asst_id> and have it just work with az login / managed identity / service principal, the same way azure/* models already do.

Problem

target.model: azure_ai/agents/<id> was technically reachable via LiteLLM, but the existing AAD wiring only injected tokens for azure/* models. Users hit two confusing dead ends:

  1. AZURE_API_KEY (Azure OpenAI) silently suppressed AAD injection for azure_ai/* calls because mode resolution was global. With a key set for orchestration (azure/gpt-4o-mini systematize/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.
  2. LiteLLM's missing-api_base error was misclassified as retryable, so the eval burned 5 retries before surfacing the real cause.

Changes

File Change
assert_ai/core/azure_auth.py Add AZURE_FOUNDRY_SCOPE. Per-scope get_azure_token_provider(scope=...) cache that shares one DefaultAzureCredential. Family-aware resolve_azure_auth_mode(family=...) that checks AZURE_API_KEY for the azure family and AZURE_AI_API_KEY for the azure_ai family. Broaden boot-log label to Azure auth mode: since it covers both families.
assert_ai/core/model_client.py New _aad_scope_for_model() dispatch. Inject AAD credentials for both Azure families: azure/* keeps azure_ad_token_provider callable, azure_ai/* calls the provider once per request and stuffs the bearer into payload['api_key'] (what LiteLLM's azure_ai/agents provider expects). Re-resolve mode per family on azure_ai/* injection to avoid the AZURE_API_KEY leak. Classifier widened so 401s and missing-token/AZURE_AI_API_BASE errors on azure_ai/* surface actionable hints instead of generic retryable errors.
assert_ai/core/config_model.py Reject target.tools and target.system_prompt on azure_ai/agents/* (the hosted agent owns its tools and instructions server-side).
docs/getting-started.md, docs/config/schema.md Two-line addendum noting Foundry agents share the same auth setup; new schema bullet for azure_ai/agents/<AGENT_ID>.
tests +8 tests: per-scope cache, family-aware resolver, env-isolation regression for the AZURE_API_KEY leak, both LiteLLM api_base error variants, hosted-agent validation guards.

Commits (6, each independently green)

SHA Subject
924eced5 feat(azure-auth): per-scope token-provider cache for azure_ai/* support
4b3f91b0 feat(model_client): inject AAD credentials for azure_ai/* models
a2699d23 feat(model_client): friendly error hints for azure_ai/* failures
4429314d chore(azure-auth): broaden boot-log label to cover azure_ai/* family
7b23b57a feat(config): reject target.tools and system_prompt on azure_ai/agents/*
1f97de8d docs(azure-auth): note azure_ai/* (Foundry agents) shares the same auth

User-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, point target.model at the agent and ASSERT mints the AAD token for you:

pipeline:
  inference:
    target:
      model:
        name: azure_ai/agents/asst_xxx
pip install -e ".[azure-aad]"
az login
export AZURE_AI_API_BASE=https://<resource>.services.ai.azure.com/api/projects/<project>
assert-ai run --config <your_config>.yaml

Service Principal (AZURE_TENANT_ID / AZURE_CLIENT_ID / AZURE_CLIENT_SECRET) and managed identity flow through the same DefaultAzureCredential chain.

Testing

  • Full suite: 1113 passed, 17 skipped, 1 deselected (docker-env test unrelated).
  • Each commit green standalone (rebase-friendly history).
  • End-to-end smoke-tested against a real Foundry-hosted Assistant: AAD token minted with https://ai.azure.com/.default scope, injected into payload['api_key'], request reached the Foundry endpoint, server-side run completed, reply returned correctly.

Out of scope

  • v2 Foundry Agents (UUID IDs at {project}/agents/{uuid}/threads/...). LiteLLM's azure_ai/agents provider only speaks the v1 Assistants surface (asst_xxx) and rejects UUIDs with Invalid '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.
  • OTel trace visibility for hosted Foundry agents. Foundry server-side traces would need a separate export integration; out of scope here.

Related

tangym added 6 commits June 25, 2026 02:15
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
tangym force-pushed the feat/foundry-agent-target branch from 1f97de8 to e0f4eb0 Compare June 25, 2026 02:16
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants