Skip to main content

agent_framework_azure/
credentials.rs

1//! A small Microsoft Entra ID (Azure AD) credential chain implementing
2//! [`TokenCredential`], with per-credential, per-scope token caching.
3//!
4//! These are hand-rolled analogues of `azure_identity`'s credentials, built on
5//! `reqwest`/`tokio` with no Azure SDK dependency:
6//!
7//! * [`AzureCliCredential`] — shells out to `az account get-access-token`.
8//! * [`ClientSecretCredential`] — OAuth2 client-credentials flow against Entra.
9//! * [`EnvironmentCredential`] — the same flow, configured from
10//!   `AZURE_TENANT_ID`/`AZURE_CLIENT_ID`/`AZURE_CLIENT_SECRET`.
11//! * [`ManagedIdentityCredential`] — the IMDS token endpoint.
12//! * [`WorkloadIdentityCredential`] — exchanges a Kubernetes federated
13//!   service-account token (AKS workload identity) for an Entra token.
14//! * [`ChainedTokenCredential`] — tries each in order; the first to succeed is
15//!   remembered and preferred thereafter.
16//! * [`DefaultAzureCredential`] — a prebuilt chain of the above four, in
17//!   `azure_identity`'s documented `DefaultAzureCredential` order.
18//!
19//! Every credential is bound to a default scope (audience) supplied at
20//! construction and used by [`TokenCredential::get_token`]; a different scope
21//! can be requested per call via [`TokenCredential::get_token_for_scope`].
22//! Tokens are cached per scope and refreshed [`REFRESH_SKEW`] before expiry.
23
24use std::collections::HashMap;
25use std::sync::{Arc, Mutex};
26use std::time::{Duration, SystemTime, UNIX_EPOCH};
27
28use agent_framework_core::error::{Error, Result};
29use async_trait::async_trait;
30use serde_json::Value;
31
32use crate::credential::TokenCredential;
33
34/// How long before a token's stated expiry it is considered stale and
35/// proactively refreshed (2 minutes), so an in-flight request never races the
36/// exact expiry instant.
37pub const REFRESH_SKEW: Duration = Duration::from_secs(120);
38
39/// The IMDS token endpoint used by [`ManagedIdentityCredential`] when no
40/// endpoint is configured or discovered from the environment.
41pub const DEFAULT_IMDS_ENDPOINT: &str = "http://169.254.169.254/metadata/identity/oauth2/token";
42
43/// The IMDS token API version.
44const IMDS_API_VERSION: &str = "2018-02-01";
45
46/// The default Entra ID authority (login endpoint) for
47/// [`ClientSecretCredential`].
48pub const DEFAULT_AUTHORITY: &str = "https://login.microsoftonline.com";
49
50/// Fallback token lifetime, in seconds, used only when a token response omits
51/// any expiry hint (1 hour is the Entra default access-token lifetime).
52const DEFAULT_TOKEN_TTL_SECS: u64 = 3600;
53
54/// Conservative lifetime, in seconds, assumed for an Azure CLI token whose
55/// JSON carries no machine-readable `expires_on` epoch (older `az` builds emit
56/// only a local-time `expiresOn` string, which cannot be converted to an
57/// instant without a timezone database). Kept short so a real refresh happens
58/// soon rather than trusting a possibly-stale token.
59const CLI_FALLBACK_TTL_SECS: u64 = 300;
60
61// ---------------------------------------------------------------------------
62// Token cache
63// ---------------------------------------------------------------------------
64
65#[derive(Clone)]
66struct CachedToken {
67    token: String,
68    expires_at: SystemTime,
69}
70
71/// A thread-safe, per-scope token cache shared by each credential.
72#[derive(Default)]
73struct TokenCache {
74    entries: Mutex<HashMap<String, CachedToken>>,
75}
76
77impl TokenCache {
78    /// Return the cached token for `scope` if one is present and not within
79    /// [`REFRESH_SKEW`] of expiry.
80    fn get(&self, scope: &str) -> Option<String> {
81        let entries = self.entries.lock().unwrap();
82        let entry = entries.get(scope)?;
83        if SystemTime::now() + REFRESH_SKEW < entry.expires_at {
84            Some(entry.token.clone())
85        } else {
86            None
87        }
88    }
89
90    /// Store `token` for `scope`, valid until `expires_at`.
91    fn put(&self, scope: &str, token: String, expires_at: SystemTime) {
92        self.entries
93            .lock()
94            .unwrap()
95            .insert(scope.to_string(), CachedToken { token, expires_at });
96    }
97}
98
99// ---------------------------------------------------------------------------
100// Shared parsing helpers
101// ---------------------------------------------------------------------------
102
103/// Read a JSON value that may be a number or a numeric string as `u64`.
104///
105/// IMDS returns `expires_in`/`expires_on` as strings; the OAuth2 token endpoint
106/// returns `expires_in` as a number — both are accepted.
107fn json_u64(v: Option<&Value>) -> Option<u64> {
108    match v {
109        Some(Value::Number(n)) => n.as_u64().or_else(|| n.as_f64().map(|f| f as u64)),
110        Some(Value::String(s)) => s.trim().parse().ok(),
111        _ => None,
112    }
113}
114
115/// Parse an OAuth2 / IMDS token response body: `access_token` plus a relative
116/// `expires_in` (seconds from now). Shared by the client-secret and
117/// managed-identity credentials.
118fn parse_oauth_token(body: &str) -> Result<(String, SystemTime)> {
119    let v: Value = serde_json::from_str(body)
120        .map_err(|e| Error::other(format!("invalid token response json: {e}")))?;
121    let token = v
122        .get("access_token")
123        .and_then(Value::as_str)
124        .ok_or_else(|| Error::other("token response missing 'access_token'"))?
125        .to_string();
126    let ttl = json_u64(v.get("expires_in")).unwrap_or(DEFAULT_TOKEN_TTL_SECS);
127    Ok((token, SystemTime::now() + Duration::from_secs(ttl)))
128}
129
130/// Parse the JSON emitted by `az account get-access-token --output json`.
131///
132/// Extracts `accessToken` and computes expiry from the integer epoch
133/// `expires_on` when present (newer `az`); otherwise falls back to a short
134/// conservative TTL (see [`CLI_FALLBACK_TTL_SECS`]).
135fn parse_cli_output(stdout: &[u8]) -> Result<(String, SystemTime)> {
136    let v: Value = serde_json::from_slice(stdout)
137        .map_err(|e| Error::other(format!("invalid Azure CLI token json: {e}")))?;
138    let token = v
139        .get("accessToken")
140        .and_then(Value::as_str)
141        .ok_or_else(|| Error::other("Azure CLI token json missing 'accessToken'"))?
142        .to_string();
143    let expires_at = match v.get("expires_on").and_then(Value::as_i64) {
144        Some(epoch) if epoch > 0 => UNIX_EPOCH + Duration::from_secs(epoch as u64),
145        _ => SystemTime::now() + Duration::from_secs(CLI_FALLBACK_TTL_SECS),
146    };
147    Ok((token, expires_at))
148}
149
150/// Convert an Entra ID scope (`"<resource>/.default"`) to the bare resource URI
151/// IMDS expects (`resource=<resource>`).
152fn resource_from_scope(scope: &str) -> &str {
153    scope.strip_suffix("/.default").unwrap_or(scope)
154}
155
156/// Names of `vars` whose resolved value is `None`, in the given order — used
157/// by [`EnvironmentCredential::new`] and [`WorkloadIdentityCredential`] to
158/// build a single "missing: X, Y" error message instead of failing on just
159/// the first missing variable.
160fn missing_env_vars<'a>(vars: &[(&'a str, &Option<String>)]) -> Vec<&'a str> {
161    vars.iter()
162        .filter(|(_, v)| v.is_none())
163        .map(|(name, _)| *name)
164        .collect()
165}
166
167// ---------------------------------------------------------------------------
168// AzureCliCredential
169// ---------------------------------------------------------------------------
170
171/// Authenticates by shelling out to the Azure CLI
172/// (`az account get-access-token --scope <scope> --output json`).
173///
174/// Useful for local development where a developer is already signed in via
175/// `az login`. Requires the `az` binary on `PATH`; a clear error is returned
176/// when it is missing.
177pub struct AzureCliCredential {
178    program: String,
179    default_scope: String,
180    cache: TokenCache,
181}
182
183impl AzureCliCredential {
184    /// Create a credential that acquires tokens for `scope` (e.g.
185    /// `"https://ai.azure.com/.default"`).
186    pub fn new(scope: impl Into<String>) -> Self {
187        Self {
188            program: "az".to_string(),
189            default_scope: scope.into(),
190            cache: TokenCache::default(),
191        }
192    }
193
194    /// Override the CLI executable (default `"az"`), e.g. an absolute path or
195    /// `"az.cmd"` on Windows.
196    pub fn with_command(mut self, program: impl Into<String>) -> Self {
197        self.program = program.into();
198        self
199    }
200
201    async fn token(&self, scope: &str) -> Result<String> {
202        if let Some(t) = self.cache.get(scope) {
203            return Ok(t);
204        }
205        let (token, expires_at) = self.fetch(scope).await?;
206        self.cache.put(scope, token.clone(), expires_at);
207        Ok(token)
208    }
209
210    async fn fetch(&self, scope: &str) -> Result<(String, SystemTime)> {
211        let output = tokio::process::Command::new(&self.program)
212            .arg("account")
213            .arg("get-access-token")
214            .arg("--scope")
215            .arg(scope)
216            .arg("--output")
217            .arg("json")
218            .output()
219            .await;
220        match output {
221            Err(e) if e.kind() == std::io::ErrorKind::NotFound => {
222                Err(Error::Configuration(format!(
223                    "Azure CLI ('{}') was not found on PATH; run `az login` after installing it, \
224                     or use a different credential: {e}",
225                    self.program
226                )))
227            }
228            Err(e) => Err(Error::service(format!("failed to run Azure CLI: {e}"))),
229            Ok(out) if !out.status.success() => {
230                let stderr = String::from_utf8_lossy(&out.stderr);
231                Err(Error::other(format!(
232                    "Azure CLI token request failed ({}): {}",
233                    out.status,
234                    stderr.trim()
235                )))
236            }
237            Ok(out) => parse_cli_output(&out.stdout),
238        }
239    }
240}
241
242#[async_trait]
243impl TokenCredential for AzureCliCredential {
244    async fn get_token(&self) -> Result<String> {
245        self.token(&self.default_scope).await
246    }
247    async fn get_token_for_scope(&self, scope: &str) -> Result<String> {
248        self.token(scope).await
249    }
250}
251
252// ---------------------------------------------------------------------------
253// ClientSecretCredential
254// ---------------------------------------------------------------------------
255
256/// Authenticates a confidential client via the OAuth2 client-credentials flow
257/// (`POST {authority}/{tenant}/oauth2/v2.0/token`).
258pub struct ClientSecretCredential {
259    http: reqwest::Client,
260    authority: String,
261    tenant_id: String,
262    client_id: String,
263    client_secret: String,
264    default_scope: String,
265    cache: TokenCache,
266}
267
268impl ClientSecretCredential {
269    /// Create a client-secret credential for the given tenant/app registration,
270    /// acquiring tokens for `scope`.
271    pub fn new(
272        tenant_id: impl Into<String>,
273        client_id: impl Into<String>,
274        client_secret: impl Into<String>,
275        scope: impl Into<String>,
276    ) -> Self {
277        Self {
278            http: reqwest::Client::new(),
279            authority: DEFAULT_AUTHORITY.to_string(),
280            tenant_id: tenant_id.into(),
281            client_id: client_id.into(),
282            client_secret: client_secret.into(),
283            default_scope: scope.into(),
284            cache: TokenCache::default(),
285        }
286    }
287
288    /// Override the Entra authority (default
289    /// [`DEFAULT_AUTHORITY`]) — e.g. a sovereign cloud, or a loopback in tests.
290    pub fn with_authority(mut self, authority: impl Into<String>) -> Self {
291        self.authority = authority.into();
292        self
293    }
294
295    fn token_url(&self) -> String {
296        format!(
297            "{}/{}/oauth2/v2.0/token",
298            self.authority.trim_end_matches('/'),
299            self.tenant_id
300        )
301    }
302
303    async fn token(&self, scope: &str) -> Result<String> {
304        if let Some(t) = self.cache.get(scope) {
305            return Ok(t);
306        }
307        let (token, expires_at) = self.fetch(scope).await?;
308        self.cache.put(scope, token.clone(), expires_at);
309        Ok(token)
310    }
311
312    async fn fetch(&self, scope: &str) -> Result<(String, SystemTime)> {
313        let params = [
314            ("grant_type", "client_credentials"),
315            ("client_id", self.client_id.as_str()),
316            ("client_secret", self.client_secret.as_str()),
317            ("scope", scope),
318        ];
319        let resp = self
320            .http
321            .post(self.token_url())
322            .form(&params)
323            .send()
324            .await
325            .map_err(|e| Error::service(format!("client-secret token request failed: {e}")))?;
326        let status = resp.status();
327        let body = resp.text().await.unwrap_or_default();
328        if !status.is_success() {
329            return Err(Error::other(format!(
330                "client-secret token request rejected ({}): {}",
331                status,
332                body.trim()
333            )));
334        }
335        parse_oauth_token(&body)
336    }
337}
338
339#[async_trait]
340impl TokenCredential for ClientSecretCredential {
341    async fn get_token(&self) -> Result<String> {
342        self.token(&self.default_scope).await
343    }
344    async fn get_token_for_scope(&self, scope: &str) -> Result<String> {
345        self.token(scope).await
346    }
347}
348
349// ---------------------------------------------------------------------------
350// EnvironmentCredential
351// ---------------------------------------------------------------------------
352
353/// Authenticates a service principal configured entirely by environment
354/// variables, via the client-secret flow: `AZURE_TENANT_ID`,
355/// `AZURE_CLIENT_ID`, `AZURE_CLIENT_SECRET`.
356///
357/// A restricted port of `azure_identity`'s `EnvironmentCredential`: upstream
358/// also supports a certificate flow and a (deprecated) username/password
359/// flow, neither of which this crate implements, so only the client-secret
360/// configuration is recognized. Unlike upstream — which defers the
361/// "not configured" failure to the first
362/// [`get_token`](TokenCredential::get_token) call — [`new`](Self::new) fails
363/// immediately, listing every missing variable, so misconfiguration is caught
364/// at startup rather than on first use.
365pub struct EnvironmentCredential {
366    inner: ClientSecretCredential,
367}
368
369impl EnvironmentCredential {
370    /// Read `AZURE_TENANT_ID`, `AZURE_CLIENT_ID`, and `AZURE_CLIENT_SECRET`
371    /// from the environment and build a client-secret credential acquiring
372    /// tokens for `scope`.
373    ///
374    /// # Errors
375    /// [`Error::Configuration`] naming every missing variable, when one or
376    /// more of the three are unset.
377    pub fn new(scope: impl Into<String>) -> Result<Self> {
378        let tenant_id = std::env::var("AZURE_TENANT_ID").ok();
379        let client_id = std::env::var("AZURE_CLIENT_ID").ok();
380        let client_secret = std::env::var("AZURE_CLIENT_SECRET").ok();
381
382        let missing = missing_env_vars(&[
383            ("AZURE_TENANT_ID", &tenant_id),
384            ("AZURE_CLIENT_ID", &client_id),
385            ("AZURE_CLIENT_SECRET", &client_secret),
386        ]);
387        if !missing.is_empty() {
388            return Err(Error::Configuration(format!(
389                "EnvironmentCredential: missing required environment variable(s): {}",
390                missing.join(", ")
391            )));
392        }
393
394        Ok(Self {
395            inner: ClientSecretCredential::new(
396                tenant_id.unwrap(),
397                client_id.unwrap(),
398                client_secret.unwrap(),
399                scope,
400            ),
401        })
402    }
403
404    /// Override the Entra authority used by the inner client-secret
405    /// credential (default [`DEFAULT_AUTHORITY`]).
406    pub fn with_authority(mut self, authority: impl Into<String>) -> Self {
407        self.inner = self.inner.with_authority(authority);
408        self
409    }
410}
411
412#[async_trait]
413impl TokenCredential for EnvironmentCredential {
414    async fn get_token(&self) -> Result<String> {
415        self.inner.get_token().await
416    }
417    async fn get_token_for_scope(&self, scope: &str) -> Result<String> {
418        self.inner.get_token_for_scope(scope).await
419    }
420}
421
422// ---------------------------------------------------------------------------
423// ManagedIdentityCredential
424// ---------------------------------------------------------------------------
425
426/// Authenticates via an Azure Managed Identity by calling the Instance Metadata
427/// Service (IMDS) token endpoint
428/// (`GET {endpoint}?api-version=2018-02-01&resource=<resource>` with the
429/// `Metadata: true` header).
430///
431/// The endpoint defaults to the IMDS address but is overridden from the
432/// `IDENTITY_ENDPOINT`/`MSI_ENDPOINT` environment variables (as set by App
433/// Service / Functions) or via [`with_endpoint`](Self::with_endpoint) — which
434/// also makes it loopback-testable. An optional user-assigned identity client
435/// id may be supplied.
436pub struct ManagedIdentityCredential {
437    http: reqwest::Client,
438    endpoint: String,
439    client_id: Option<String>,
440    identity_header: Option<String>,
441    default_scope: String,
442    cache: TokenCache,
443}
444
445impl ManagedIdentityCredential {
446    /// Create a managed-identity credential acquiring tokens for `scope`.
447    ///
448    /// The IMDS endpoint is taken from `IDENTITY_ENDPOINT`, then `MSI_ENDPOINT`,
449    /// then [`DEFAULT_IMDS_ENDPOINT`]; an `IDENTITY_HEADER`, when present, is
450    /// forwarded as `X-IDENTITY-HEADER`.
451    pub fn new(scope: impl Into<String>) -> Self {
452        let endpoint = std::env::var("IDENTITY_ENDPOINT")
453            .or_else(|_| std::env::var("MSI_ENDPOINT"))
454            .unwrap_or_else(|_| DEFAULT_IMDS_ENDPOINT.to_string());
455        let identity_header = std::env::var("IDENTITY_HEADER").ok();
456        // IMDS / App Service MSI endpoints are link-local and answer within
457        // milliseconds when present. Bound the request so a machine with no
458        // managed identity (e.g. local dev) fails this chain link quickly,
459        // instead of stalling `DefaultAzureCredential` for the OS TCP
460        // timeout before it can fall through to `AzureCliCredential`.
461        let http = reqwest::Client::builder()
462            .connect_timeout(std::time::Duration::from_secs(2))
463            .timeout(std::time::Duration::from_secs(10))
464            .build()
465            .unwrap_or_else(|_| reqwest::Client::new());
466        Self {
467            http,
468            endpoint,
469            client_id: None,
470            identity_header,
471            default_scope: scope.into(),
472            cache: TokenCache::default(),
473        }
474    }
475
476    /// Override the token endpoint (default [`DEFAULT_IMDS_ENDPOINT`] or the
477    /// environment) — the seam used by loopback tests.
478    pub fn with_endpoint(mut self, endpoint: impl Into<String>) -> Self {
479        self.endpoint = endpoint.into();
480        self
481    }
482
483    /// Pin a user-assigned managed identity by its client id.
484    pub fn with_client_id(mut self, client_id: impl Into<String>) -> Self {
485        self.client_id = Some(client_id.into());
486        self
487    }
488
489    /// Set the `X-IDENTITY-HEADER` value (App Service / Functions secret).
490    pub fn with_identity_header(mut self, header: impl Into<String>) -> Self {
491        self.identity_header = Some(header.into());
492        self
493    }
494
495    async fn token(&self, scope: &str) -> Result<String> {
496        if let Some(t) = self.cache.get(scope) {
497            return Ok(t);
498        }
499        let (token, expires_at) = self.fetch(scope).await?;
500        self.cache.put(scope, token.clone(), expires_at);
501        Ok(token)
502    }
503
504    async fn fetch(&self, scope: &str) -> Result<(String, SystemTime)> {
505        let resource = resource_from_scope(scope);
506        let mut params: Vec<(&str, &str)> =
507            vec![("api-version", IMDS_API_VERSION), ("resource", resource)];
508        if let Some(cid) = &self.client_id {
509            params.push(("client_id", cid.as_str()));
510        }
511        let mut req = self
512            .http
513            .get(&self.endpoint)
514            .query(&params)
515            .header("Metadata", "true");
516        if let Some(h) = &self.identity_header {
517            req = req.header("X-IDENTITY-HEADER", h);
518        }
519        let resp = req
520            .send()
521            .await
522            .map_err(|e| Error::service(format!("managed-identity token request failed: {e}")))?;
523        let status = resp.status();
524        let body = resp.text().await.unwrap_or_default();
525        if !status.is_success() {
526            return Err(Error::other(format!(
527                "managed-identity token request rejected ({}): {}",
528                status,
529                body.trim()
530            )));
531        }
532        parse_oauth_token(&body)
533    }
534}
535
536#[async_trait]
537impl TokenCredential for ManagedIdentityCredential {
538    async fn get_token(&self) -> Result<String> {
539        self.token(&self.default_scope).await
540    }
541    async fn get_token_for_scope(&self, scope: &str) -> Result<String> {
542        self.token(scope).await
543    }
544}
545
546// ---------------------------------------------------------------------------
547// WorkloadIdentityCredential
548// ---------------------------------------------------------------------------
549
550/// Authenticates via Microsoft Entra Workload ID (Azure Kubernetes Service
551/// workload identity federation): exchanges a Kubernetes projected
552/// service-account token for an Entra access token using the OAuth2
553/// client-credentials grant with a JWT client assertion
554/// (`POST {authority}/{tenant}/oauth2/v2.0/token`,
555/// `client_assertion_type=urn:ietf:params:oauth:client-assertion-type:jwt-bearer`).
556///
557/// By default `tenant_id`, `client_id`, and the federated token file path are
558/// read from `AZURE_TENANT_ID`, `AZURE_CLIENT_ID`, and
559/// `AZURE_FEDERATED_TOKEN_FILE` — the variables the AKS workload identity
560/// webhook injects into a pod — mirroring `azure_identity`'s
561/// `WorkloadIdentityCredential`. The token file is re-read on every token
562/// *request* (not cached once at startup), since the kubelet periodically
563/// rotates the projected token.
564pub struct WorkloadIdentityCredential {
565    http: reqwest::Client,
566    authority: String,
567    tenant_id: String,
568    client_id: String,
569    token_file_path: String,
570    default_scope: String,
571    cache: TokenCache,
572}
573
574impl WorkloadIdentityCredential {
575    /// Read `AZURE_TENANT_ID`, `AZURE_CLIENT_ID`, and
576    /// `AZURE_FEDERATED_TOKEN_FILE` from the environment, acquiring tokens for
577    /// `scope`.
578    ///
579    /// # Errors
580    /// [`Error::Configuration`] naming every missing variable, when one or
581    /// more of the three are unset.
582    pub fn new(scope: impl Into<String>) -> Result<Self> {
583        Self::with_overrides(scope, None, None, None)
584    }
585
586    /// Same as [`new`](Self::new), but the tenant id, client id, and/or
587    /// federated-token-file path can be supplied directly instead of read
588    /// from the environment. Pass `None` for a field to keep reading it from
589    /// the usual environment variable.
590    ///
591    /// # Errors
592    /// [`Error::Configuration`] naming every value that is neither passed
593    /// here nor found in the environment.
594    pub fn with_overrides(
595        scope: impl Into<String>,
596        tenant_id: Option<String>,
597        client_id: Option<String>,
598        token_file_path: Option<String>,
599    ) -> Result<Self> {
600        let tenant_id = tenant_id.or_else(|| std::env::var("AZURE_TENANT_ID").ok());
601        let client_id = client_id.or_else(|| std::env::var("AZURE_CLIENT_ID").ok());
602        let token_file_path =
603            token_file_path.or_else(|| std::env::var("AZURE_FEDERATED_TOKEN_FILE").ok());
604
605        let missing = missing_env_vars(&[
606            ("AZURE_TENANT_ID", &tenant_id),
607            ("AZURE_CLIENT_ID", &client_id),
608            ("AZURE_FEDERATED_TOKEN_FILE", &token_file_path),
609        ]);
610        if !missing.is_empty() {
611            return Err(Error::Configuration(format!(
612                "WorkloadIdentityCredential: missing required configuration: {}",
613                missing.join(", ")
614            )));
615        }
616
617        Ok(Self {
618            http: reqwest::Client::new(),
619            authority: DEFAULT_AUTHORITY.to_string(),
620            tenant_id: tenant_id.unwrap(),
621            client_id: client_id.unwrap(),
622            token_file_path: token_file_path.unwrap(),
623            default_scope: scope.into(),
624            cache: TokenCache::default(),
625        })
626    }
627
628    /// Same as [`new`](Self::new), overriding only the client id (default:
629    /// `AZURE_CLIENT_ID`) — e.g. to authenticate as an app registration whose
630    /// federated credential trusts this pod's service account token, distinct
631    /// from the identity the AKS webhook injects by default.
632    pub fn with_client_id(scope: impl Into<String>, client_id: Option<String>) -> Result<Self> {
633        Self::with_overrides(scope, None, client_id, None)
634    }
635
636    /// Override the Entra authority (default [`DEFAULT_AUTHORITY`]) — e.g. a
637    /// sovereign cloud, or a loopback in tests.
638    pub fn with_authority(mut self, authority: impl Into<String>) -> Self {
639        self.authority = authority.into();
640        self
641    }
642
643    fn token_url(&self) -> String {
644        format!(
645            "{}/{}/oauth2/v2.0/token",
646            self.authority.trim_end_matches('/'),
647            self.tenant_id
648        )
649    }
650
651    async fn token(&self, scope: &str) -> Result<String> {
652        if let Some(t) = self.cache.get(scope) {
653            return Ok(t);
654        }
655        let (token, expires_at) = self.fetch(scope).await?;
656        self.cache.put(scope, token.clone(), expires_at);
657        Ok(token)
658    }
659
660    /// Build the token-request form body, re-reading the federated token file
661    /// fresh every time so a rotated projected token is always picked up.
662    /// Factored out from [`fetch`](Self::fetch) so the exact request shape is
663    /// unit-testable without a network.
664    fn request_params(&self, scope: &str) -> Result<Vec<(String, String)>> {
665        let assertion = std::fs::read_to_string(&self.token_file_path).map_err(|e| {
666            Error::Configuration(format!(
667                "WorkloadIdentityCredential: failed to read federated token file '{}': {e}",
668                self.token_file_path
669            ))
670        })?;
671        Ok(vec![
672            ("grant_type".to_string(), "client_credentials".to_string()),
673            ("client_id".to_string(), self.client_id.clone()),
674            ("client_assertion".to_string(), assertion.trim().to_string()),
675            (
676                "client_assertion_type".to_string(),
677                "urn:ietf:params:oauth:client-assertion-type:jwt-bearer".to_string(),
678            ),
679            ("scope".to_string(), scope.to_string()),
680        ])
681    }
682
683    async fn fetch(&self, scope: &str) -> Result<(String, SystemTime)> {
684        let params = self.request_params(scope)?;
685        let resp = self
686            .http
687            .post(self.token_url())
688            .form(&params)
689            .send()
690            .await
691            .map_err(|e| Error::service(format!("workload-identity token request failed: {e}")))?;
692        let status = resp.status();
693        let body = resp.text().await.unwrap_or_default();
694        if !status.is_success() {
695            return Err(Error::other(format!(
696                "workload-identity token request rejected ({}): {}",
697                status,
698                body.trim()
699            )));
700        }
701        parse_oauth_token(&body)
702    }
703}
704
705#[async_trait]
706impl TokenCredential for WorkloadIdentityCredential {
707    async fn get_token(&self) -> Result<String> {
708        self.token(&self.default_scope).await
709    }
710    async fn get_token_for_scope(&self, scope: &str) -> Result<String> {
711        self.token(scope).await
712    }
713}
714
715// ---------------------------------------------------------------------------
716// ChainedTokenCredential
717// ---------------------------------------------------------------------------
718
719/// Tries a sequence of credentials in order, returning the first token
720/// obtained. The index of the first credential to succeed is remembered and
721/// tried first on subsequent calls (re-remembered if a different one later
722/// wins), mirroring `azure_identity`'s `ChainedTokenCredential`.
723pub struct ChainedTokenCredential {
724    credentials: Vec<Arc<dyn TokenCredential>>,
725    remembered: Mutex<Option<usize>>,
726}
727
728impl ChainedTokenCredential {
729    /// Build a chain from an ordered list of credentials.
730    pub fn new(credentials: Vec<Arc<dyn TokenCredential>>) -> Self {
731        Self {
732            credentials,
733            remembered: Mutex::new(None),
734        }
735    }
736
737    async fn acquire(&self, scope: Option<&str>) -> Result<String> {
738        // Copy the remembered index out before awaiting; never hold the lock
739        // across an `.await`.
740        let remembered = *self.remembered.lock().unwrap();
741
742        if let Some(idx) = remembered {
743            if let Some(cred) = self.credentials.get(idx) {
744                if let Ok(token) = Self::call(cred.as_ref(), scope).await {
745                    return Ok(token);
746                }
747            }
748        }
749
750        let mut errors = Vec::new();
751        for (i, cred) in self.credentials.iter().enumerate() {
752            if Some(i) == remembered {
753                // Already attempted above.
754                continue;
755            }
756            match Self::call(cred.as_ref(), scope).await {
757                Ok(token) => {
758                    *self.remembered.lock().unwrap() = Some(i);
759                    return Ok(token);
760                }
761                Err(e) => errors.push(format!("credential[{i}]: {e}")),
762            }
763        }
764
765        Err(Error::Configuration(format!(
766            "ChainedTokenCredential: no credential in the chain returned a token ({})",
767            if errors.is_empty() {
768                "the chain is empty".to_string()
769            } else {
770                errors.join("; ")
771            }
772        )))
773    }
774
775    async fn call(cred: &dyn TokenCredential, scope: Option<&str>) -> Result<String> {
776        match scope {
777            Some(s) => cred.get_token_for_scope(s).await,
778            None => cred.get_token().await,
779        }
780    }
781}
782
783#[async_trait]
784impl TokenCredential for ChainedTokenCredential {
785    async fn get_token(&self) -> Result<String> {
786        self.acquire(None).await
787    }
788    async fn get_token_for_scope(&self, scope: &str) -> Result<String> {
789        self.acquire(Some(scope)).await
790    }
791}
792
793// ---------------------------------------------------------------------------
794// DefaultAzureCredential
795// ---------------------------------------------------------------------------
796
797/// A prebuilt [`ChainedTokenCredential`], in (a subset of) `azure_identity`'s
798/// documented `DefaultAzureCredential` order.
799///
800/// Upstream (Python/.NET) `DefaultAzureCredential` tries, in order:
801/// `EnvironmentCredential`, `WorkloadIdentityCredential`,
802/// `ManagedIdentityCredential`, `SharedTokenCacheCredential` (Windows only),
803/// `VisualStudioCodeCredential`, `AzureCliCredential`,
804/// `AzurePowerShellCredential`, `AzureDeveloperCliCredential`, and
805/// (Windows/WSL only, off by default) `InteractiveBrowserCredential`/broker
806/// authentication.
807///
808/// This crate implements four of those credential types, so this chain is
809/// **`EnvironmentCredential` → `WorkloadIdentityCredential` →
810/// `ManagedIdentityCredential` → `AzureCliCredential`** — the same relative
811/// order upstream uses among the credentials that exist here. Every other
812/// link in upstream's list (the shared token cache, IDE-signed-in
813/// credentials, PowerShell, the Azure Developer CLI, interactive/broker auth)
814/// is simply absent rather than approximated by something else; callers that
815/// need one of those should build their own [`ChainedTokenCredential`].
816///
817/// `EnvironmentCredential` and `WorkloadIdentityCredential` are included only
818/// when their required environment variables are present — construction
819/// failure just excludes them from the chain, mirroring upstream's
820/// "unavailable credentials are skipped, not fatal" behavior (there, via a
821/// placeholder that defers the error to `get_token`; here, by omission, since
822/// [`ChainedTokenCredential`] already only holds constructed credentials).
823/// `ManagedIdentityCredential` and `AzureCliCredential` need no such
824/// configuration and are always present, so the chain is never empty. As with
825/// any [`ChainedTokenCredential`], whichever credential first returns a token
826/// is remembered and tried first on subsequent calls.
827pub struct DefaultAzureCredential {
828    chain: ChainedTokenCredential,
829}
830
831impl DefaultAzureCredential {
832    /// Build the chain, acquiring tokens for `scope` by default.
833    pub fn new(scope: impl Into<String>) -> Self {
834        let scope = scope.into();
835        Self {
836            chain: ChainedTokenCredential::new(default_chain(&scope)),
837        }
838    }
839}
840
841/// The ordered, environment-filtered credential list backing
842/// [`DefaultAzureCredential`]. A free function so the inclusion logic is
843/// independently unit-testable without a network.
844fn default_chain(scope: &str) -> Vec<Arc<dyn TokenCredential>> {
845    let mut chain: Vec<Arc<dyn TokenCredential>> = Vec::new();
846    if let Ok(cred) = EnvironmentCredential::new(scope) {
847        chain.push(Arc::new(cred));
848    }
849    if let Ok(cred) = WorkloadIdentityCredential::new(scope) {
850        chain.push(Arc::new(cred));
851    }
852    chain.push(Arc::new(ManagedIdentityCredential::new(scope)));
853    chain.push(Arc::new(AzureCliCredential::new(scope)));
854    chain
855}
856
857#[async_trait]
858impl TokenCredential for DefaultAzureCredential {
859    async fn get_token(&self) -> Result<String> {
860        self.chain.get_token().await
861    }
862    async fn get_token_for_scope(&self, scope: &str) -> Result<String> {
863        self.chain.get_token_for_scope(scope).await
864    }
865}
866
867#[cfg(test)]
868mod tests {
869    use super::*;
870    use std::sync::atomic::{AtomicUsize, Ordering};
871
872    // region: token cache
873
874    #[test]
875    fn cache_returns_token_before_skew_window() {
876        let cache = TokenCache::default();
877        cache.put(
878            "scope",
879            "tok".into(),
880            SystemTime::now() + Duration::from_secs(3600),
881        );
882        assert_eq!(cache.get("scope").as_deref(), Some("tok"));
883    }
884
885    #[test]
886    fn cache_misses_within_refresh_skew() {
887        let cache = TokenCache::default();
888        // Expires in 60s, inside the 120s skew → treated as stale.
889        cache.put(
890            "scope",
891            "tok".into(),
892            SystemTime::now() + Duration::from_secs(60),
893        );
894        assert!(cache.get("scope").is_none());
895    }
896
897    #[test]
898    fn cache_is_per_scope() {
899        let cache = TokenCache::default();
900        cache.put(
901            "a",
902            "tok-a".into(),
903            SystemTime::now() + Duration::from_secs(3600),
904        );
905        assert_eq!(cache.get("a").as_deref(), Some("tok-a"));
906        assert!(cache.get("b").is_none());
907    }
908
909    // endregion
910
911    // region: parsing helpers
912
913    #[test]
914    fn parse_cli_output_reads_token_and_epoch_expiry() {
915        let json = br#"{"accessToken":"cli-token","expiresOn":"2099-01-01 00:00:00.000000","expires_on":4070908800,"tokenType":"Bearer"}"#;
916        let (token, expires_at) = parse_cli_output(json).unwrap();
917        assert_eq!(token, "cli-token");
918        assert_eq!(expires_at, UNIX_EPOCH + Duration::from_secs(4070908800));
919    }
920
921    #[test]
922    fn parse_cli_output_falls_back_to_ttl_without_epoch() {
923        let before = SystemTime::now();
924        let (token, expires_at) = parse_cli_output(br#"{"accessToken":"t"}"#).unwrap();
925        assert_eq!(token, "t");
926        assert!(expires_at >= before + Duration::from_secs(CLI_FALLBACK_TTL_SECS));
927    }
928
929    #[test]
930    fn parse_cli_output_errors_without_token() {
931        assert!(parse_cli_output(br#"{"tokenType":"Bearer"}"#).is_err());
932    }
933
934    #[test]
935    fn parse_oauth_token_number_and_string_expiry() {
936        let (t, _) = parse_oauth_token(r#"{"access_token":"a","expires_in":3600}"#).unwrap();
937        assert_eq!(t, "a");
938        // IMDS returns expires_in as a string.
939        let (t2, _) = parse_oauth_token(r#"{"access_token":"b","expires_in":"3600"}"#).unwrap();
940        assert_eq!(t2, "b");
941    }
942
943    #[test]
944    fn resource_strips_default_suffix() {
945        assert_eq!(
946            resource_from_scope("https://ai.azure.com/.default"),
947            "https://ai.azure.com"
948        );
949        assert_eq!(
950            resource_from_scope("https://vault.azure.net"),
951            "https://vault.azure.net"
952        );
953    }
954
955    // endregion
956
957    // region: azure cli missing-binary path
958
959    #[tokio::test]
960    async fn azure_cli_missing_binary_gives_clear_error() {
961        let cred = AzureCliCredential::new("https://ai.azure.com/.default")
962            .with_command("definitely-not-a-real-az-binary-xyz");
963        let err = cred.get_token().await.unwrap_err();
964        let msg = err.to_string();
965        assert!(
966            msg.contains("was not found on PATH"),
967            "expected a clear missing-CLI message, got: {msg}"
968        );
969    }
970
971    // endregion
972
973    // region: chain order + remembering
974
975    struct FakeCredential {
976        result: std::result::Result<String, ()>,
977        calls: AtomicUsize,
978    }
979
980    impl FakeCredential {
981        fn ok(token: &str) -> Self {
982            Self {
983                result: Ok(token.to_string()),
984                calls: AtomicUsize::new(0),
985            }
986        }
987        fn fail() -> Self {
988            Self {
989                result: Err(()),
990                calls: AtomicUsize::new(0),
991            }
992        }
993    }
994
995    #[async_trait]
996    impl TokenCredential for FakeCredential {
997        async fn get_token(&self) -> Result<String> {
998            self.calls.fetch_add(1, Ordering::SeqCst);
999            self.result
1000                .clone()
1001                .map_err(|_| Error::other("fake credential failure"))
1002        }
1003    }
1004
1005    #[tokio::test]
1006    async fn chain_returns_first_success_and_remembers_it() {
1007        let first = Arc::new(FakeCredential::fail());
1008        let second = Arc::new(FakeCredential::ok("winner"));
1009        let third = Arc::new(FakeCredential::ok("unused"));
1010        let chain = ChainedTokenCredential::new(vec![first.clone(), second.clone(), third.clone()]);
1011
1012        assert_eq!(chain.get_token().await.unwrap(), "winner");
1013        // First and second were tried; third never reached.
1014        assert_eq!(first.calls.load(Ordering::SeqCst), 1);
1015        assert_eq!(second.calls.load(Ordering::SeqCst), 1);
1016        assert_eq!(third.calls.load(Ordering::SeqCst), 0);
1017
1018        // Second is remembered: the next call goes straight to it, skipping the
1019        // failing first credential.
1020        assert_eq!(chain.get_token().await.unwrap(), "winner");
1021        assert_eq!(first.calls.load(Ordering::SeqCst), 1);
1022        assert_eq!(second.calls.load(Ordering::SeqCst), 2);
1023    }
1024
1025    #[tokio::test]
1026    async fn chain_reports_all_failures_when_none_succeed() {
1027        let chain = ChainedTokenCredential::new(vec![
1028            Arc::new(FakeCredential::fail()),
1029            Arc::new(FakeCredential::fail()),
1030        ]);
1031        let err = chain.get_token().await.unwrap_err();
1032        assert!(err.to_string().contains("no credential in the chain"));
1033    }
1034
1035    #[tokio::test]
1036    async fn empty_chain_errors() {
1037        let chain = ChainedTokenCredential::new(vec![]);
1038        assert!(chain.get_token().await.is_err());
1039    }
1040
1041    // endregion
1042
1043    // region: env-var-backed credentials (Environment / WorkloadIdentity / Default)
1044
1045    /// Guards mutation of the Entra env vars consumed by
1046    /// `EnvironmentCredential`/`WorkloadIdentityCredential`/
1047    /// `DefaultAzureCredential`: tests in this module run on multiple
1048    /// threads, and env vars are process-global (same pattern as
1049    /// `agent-framework-mem0`'s `ENV_MUTEX`).
1050    static ENV_MUTEX: std::sync::Mutex<()> = std::sync::Mutex::new(());
1051
1052    const ENTRA_ENV_VARS: &[&str] = &[
1053        "AZURE_TENANT_ID",
1054        "AZURE_CLIENT_ID",
1055        "AZURE_CLIENT_SECRET",
1056        "AZURE_FEDERATED_TOKEN_FILE",
1057    ];
1058
1059    fn clear_entra_env() {
1060        for key in ENTRA_ENV_VARS {
1061            // SAFETY: serialized by ENV_MUTEX; no other test in this crate
1062            // touches these variables.
1063            unsafe { std::env::remove_var(key) };
1064        }
1065    }
1066
1067    const TEST_SCOPE: &str = "https://ai.azure.com/.default";
1068
1069    #[test]
1070    fn environment_credential_errors_listing_all_missing_vars() {
1071        let _guard = ENV_MUTEX.lock().unwrap();
1072        clear_entra_env();
1073        // `.err().unwrap()` rather than `.unwrap_err()`: the credential types
1074        // deliberately don't derive `Debug` (they can hold secrets), and only
1075        // `unwrap_err()` requires the `Ok` side to be `Debug`.
1076        let err = EnvironmentCredential::new(TEST_SCOPE).err().unwrap();
1077        clear_entra_env();
1078        let msg = err.to_string();
1079        assert!(msg.contains("AZURE_TENANT_ID"), "{msg}");
1080        assert!(msg.contains("AZURE_CLIENT_ID"), "{msg}");
1081        assert!(msg.contains("AZURE_CLIENT_SECRET"), "{msg}");
1082    }
1083
1084    #[test]
1085    fn environment_credential_lists_only_the_vars_still_missing() {
1086        let _guard = ENV_MUTEX.lock().unwrap();
1087        clear_entra_env();
1088        unsafe {
1089            std::env::set_var("AZURE_TENANT_ID", "tenant");
1090            std::env::set_var("AZURE_CLIENT_ID", "client");
1091        }
1092        let err = EnvironmentCredential::new(TEST_SCOPE).err().unwrap();
1093        clear_entra_env();
1094        let msg = err.to_string();
1095        assert!(msg.contains("AZURE_CLIENT_SECRET"), "{msg}");
1096        assert!(!msg.contains("AZURE_TENANT_ID"), "{msg}");
1097        assert!(!msg.contains("AZURE_CLIENT_ID"), "{msg}");
1098    }
1099
1100    #[test]
1101    fn environment_credential_succeeds_with_all_three_vars() {
1102        let _guard = ENV_MUTEX.lock().unwrap();
1103        clear_entra_env();
1104        unsafe {
1105            std::env::set_var("AZURE_TENANT_ID", "tenant");
1106            std::env::set_var("AZURE_CLIENT_ID", "client");
1107            std::env::set_var("AZURE_CLIENT_SECRET", "secret");
1108        }
1109        let cred = EnvironmentCredential::new(TEST_SCOPE);
1110        clear_entra_env();
1111        assert!(cred.is_ok());
1112    }
1113
1114    #[test]
1115    fn workload_identity_errors_listing_all_missing_vars() {
1116        let _guard = ENV_MUTEX.lock().unwrap();
1117        clear_entra_env();
1118        let err = WorkloadIdentityCredential::new(TEST_SCOPE).err().unwrap();
1119        clear_entra_env();
1120        let msg = err.to_string();
1121        assert!(msg.contains("AZURE_TENANT_ID"), "{msg}");
1122        assert!(msg.contains("AZURE_CLIENT_ID"), "{msg}");
1123        assert!(msg.contains("AZURE_FEDERATED_TOKEN_FILE"), "{msg}");
1124    }
1125
1126    #[test]
1127    fn workload_identity_client_id_override_fills_in_for_missing_env_var() {
1128        let _guard = ENV_MUTEX.lock().unwrap();
1129        clear_entra_env();
1130        unsafe {
1131            std::env::set_var("AZURE_TENANT_ID", "tenant");
1132            std::env::set_var("AZURE_FEDERATED_TOKEN_FILE", "placeholder-token-file");
1133        }
1134        let cred = WorkloadIdentityCredential::with_client_id(
1135            TEST_SCOPE,
1136            Some("explicit-client".to_string()),
1137        );
1138        clear_entra_env();
1139        assert!(cred.is_ok());
1140    }
1141
1142    #[test]
1143    fn workload_identity_request_params_build_client_assertion_shape() {
1144        let _guard = ENV_MUTEX.lock().unwrap();
1145        // A real scratch file: `request_params` reads it fresh every call.
1146        let path = std::env::temp_dir().join(format!(
1147            "af-workload-identity-test-{}-{:?}.jwt",
1148            std::process::id(),
1149            std::thread::current().id()
1150        ));
1151        std::fs::write(&path, "fake.jwt.token").unwrap();
1152
1153        let cred = WorkloadIdentityCredential::with_overrides(
1154            TEST_SCOPE,
1155            Some("tenant-1".to_string()),
1156            Some("client-1".to_string()),
1157            Some(path.to_str().unwrap().to_string()),
1158        )
1159        .unwrap();
1160
1161        let params = cred.request_params(TEST_SCOPE).unwrap();
1162        let get = |key: &str| {
1163            params
1164                .iter()
1165                .find(|(k, _)| k == key)
1166                .map(|(_, v)| v.clone())
1167        };
1168        assert_eq!(get("grant_type").as_deref(), Some("client_credentials"));
1169        assert_eq!(get("client_id").as_deref(), Some("client-1"));
1170        assert_eq!(get("client_assertion").as_deref(), Some("fake.jwt.token"));
1171        assert_eq!(
1172            get("client_assertion_type").as_deref(),
1173            Some("urn:ietf:params:oauth:client-assertion-type:jwt-bearer")
1174        );
1175        assert_eq!(get("scope").as_deref(), Some(TEST_SCOPE));
1176        assert_eq!(
1177            cred.token_url(),
1178            "https://login.microsoftonline.com/tenant-1/oauth2/v2.0/token"
1179        );
1180
1181        std::fs::remove_file(&path).ok();
1182    }
1183
1184    #[test]
1185    fn workload_identity_rereads_token_file_on_each_request() {
1186        let path = std::env::temp_dir().join(format!(
1187            "af-workload-identity-rotate-{}-{:?}.jwt",
1188            std::process::id(),
1189            std::thread::current().id()
1190        ));
1191        std::fs::write(&path, "first-token").unwrap();
1192
1193        let cred = WorkloadIdentityCredential::with_overrides(
1194            TEST_SCOPE,
1195            Some("tenant-1".to_string()),
1196            Some("client-1".to_string()),
1197            Some(path.to_str().unwrap().to_string()),
1198        )
1199        .unwrap();
1200
1201        let first = cred.request_params("scope").unwrap();
1202        let assertion_of = |params: &[(String, String)]| {
1203            params
1204                .iter()
1205                .find(|(k, _)| k == "client_assertion")
1206                .map(|(_, v)| v.clone())
1207        };
1208        assert_eq!(assertion_of(&first).as_deref(), Some("first-token"));
1209
1210        std::fs::write(&path, "rotated-token").unwrap();
1211        let second = cred.request_params("scope").unwrap();
1212        assert_eq!(assertion_of(&second).as_deref(), Some("rotated-token"));
1213
1214        std::fs::remove_file(&path).ok();
1215    }
1216
1217    #[test]
1218    fn workload_identity_missing_token_file_is_a_clear_error() {
1219        let cred = WorkloadIdentityCredential::with_overrides(
1220            TEST_SCOPE,
1221            Some("tenant-1".to_string()),
1222            Some("client-1".to_string()),
1223            Some("/definitely/not/a/real/path/af-test.jwt".to_string()),
1224        )
1225        .unwrap();
1226        let err = cred.request_params(TEST_SCOPE).unwrap_err();
1227        assert!(err.to_string().contains("federated token file"), "{err}");
1228    }
1229
1230    // endregion
1231
1232    // region: DefaultAzureCredential chain composition/order
1233
1234    #[test]
1235    fn default_chain_includes_only_always_on_credentials_without_env() {
1236        let _guard = ENV_MUTEX.lock().unwrap();
1237        clear_entra_env();
1238        let chain = default_chain(TEST_SCOPE);
1239        clear_entra_env();
1240        // Only ManagedIdentityCredential + AzureCliCredential need no
1241        // configuration, so they're the only two present.
1242        assert_eq!(chain.len(), 2);
1243    }
1244
1245    #[test]
1246    fn default_chain_adds_environment_credential_when_configured() {
1247        let _guard = ENV_MUTEX.lock().unwrap();
1248        clear_entra_env();
1249        unsafe {
1250            std::env::set_var("AZURE_TENANT_ID", "tenant");
1251            std::env::set_var("AZURE_CLIENT_ID", "client");
1252            std::env::set_var("AZURE_CLIENT_SECRET", "secret");
1253        }
1254        let chain = default_chain(TEST_SCOPE);
1255        clear_entra_env();
1256        assert_eq!(chain.len(), 3, "Environment + ManagedIdentity + AzureCli");
1257    }
1258
1259    #[test]
1260    fn default_chain_adds_workload_identity_when_configured() {
1261        let _guard = ENV_MUTEX.lock().unwrap();
1262        clear_entra_env();
1263        unsafe {
1264            std::env::set_var("AZURE_TENANT_ID", "tenant");
1265            std::env::set_var("AZURE_CLIENT_ID", "client");
1266            std::env::set_var("AZURE_FEDERATED_TOKEN_FILE", "placeholder-token-file");
1267        }
1268        let chain = default_chain(TEST_SCOPE);
1269        clear_entra_env();
1270        assert_eq!(
1271            chain.len(),
1272            3,
1273            "WorkloadIdentity + ManagedIdentity + AzureCli"
1274        );
1275    }
1276
1277    #[test]
1278    fn default_chain_adds_both_environment_and_workload_identity_in_order() {
1279        let _guard = ENV_MUTEX.lock().unwrap();
1280        clear_entra_env();
1281        unsafe {
1282            std::env::set_var("AZURE_TENANT_ID", "tenant");
1283            std::env::set_var("AZURE_CLIENT_ID", "client");
1284            std::env::set_var("AZURE_CLIENT_SECRET", "secret");
1285            std::env::set_var("AZURE_FEDERATED_TOKEN_FILE", "placeholder-token-file");
1286        }
1287        let chain = default_chain(TEST_SCOPE);
1288        clear_entra_env();
1289        assert_eq!(chain.len(), 4, "all four credentials configured");
1290    }
1291
1292    #[test]
1293    fn default_azure_credential_wraps_a_never_empty_chain() {
1294        let _guard = ENV_MUTEX.lock().unwrap();
1295        clear_entra_env();
1296        let cred = DefaultAzureCredential::new(TEST_SCOPE);
1297        clear_entra_env();
1298        assert_eq!(cred.chain.credentials.len(), 2);
1299    }
1300
1301    // endregion
1302}