feat(azure): support Entra ID bearer token auth via AZURE_OPENAI_AD_TOKEN (#9716)

Signed-off-by: hammadxcm <hammadkhanxcm@gmail.com>
Signed-off-by: Douwe M Osinga <douwe@sidewalklabs.com>
Co-authored-by: Douwe M Osinga <douwe@sidewalklabs.com>

crates/goose-providers/src/canonical/catalog.rs

@@ -693,6 +693,12 @@ const SETUP_METADATA: &[CuratedSetupMetadata] = &[
                 placeholder: Some("Paste your API key"),
                 default_value: None,
             },
+            CuratedFieldMetadata {
+                key: "AZURE_OPENAI_AD_TOKEN",
+                label: "Entra ID Token",
+                placeholder: Some("Optional: short-lived Microsoft Entra access token"),
+                default_value: None,
+            },
         ],
     },
     CuratedSetupMetadata {

crates/goose/src/providers/azure.rs

@@ -41,7 +41,8 @@ impl AuthProvider for AzureAuthProvider {
             super::azureauth::AzureCredentials::ApiKey(_) => {
                 Ok(("api-key".to_string(), auth_token.token_value))
             }
-            super::azureauth::AzureCredentials::DefaultCredential => Ok((
+            super::azureauth::AzureCredentials::BearerToken(_)
+            | super::azureauth::AzureCredentials::DefaultCredential => Ok((
                 "Authorization".to_string(),
                 format!("Bearer {}", auth_token.token_value),
             )),
@@ -56,7 +57,7 @@ impl ProviderDef for AzureProvider {
         ProviderMetadata::new(
             AZURE_PROVIDER_NAME,
             "Azure OpenAI",
-            "Models through Azure OpenAI Service (uses Azure credential chain by default)",
+            "Models through Azure OpenAI Service (supports API key, Entra ID bearer token, and Azure credential chain)",
             "gpt-4o",
             AZURE_OPENAI_KNOWN_MODELS.to_vec(),
             AZURE_DOC_URL,
@@ -65,6 +66,7 @@ impl ProviderDef for AzureProvider {
                 ConfigKey::new("AZURE_OPENAI_DEPLOYMENT_NAME", true, false, None, true),
                 ConfigKey::new("AZURE_OPENAI_API_VERSION", false, false, None, false),
                 ConfigKey::new("AZURE_OPENAI_API_KEY", false, true, Some(""), true),
+                ConfigKey::new("AZURE_OPENAI_AD_TOKEN", false, true, Some(""), true),
             ],
         )
     }
@@ -92,7 +94,11 @@ impl ProviderDef for AzureProvider {
                 .get_secret("AZURE_OPENAI_API_KEY")
                 .ok()
                 .filter(|key: &String| !key.is_empty());
-            let auth = AzureAuth::new(api_key).map_err(|e| match e {
+            let ad_token = config
+                .get_secret("AZURE_OPENAI_AD_TOKEN")
+                .ok()
+                .filter(|token: &String| !token.is_empty());
+            let auth = AzureAuth::new(api_key, ad_token).map_err(|e| match e {
                 AuthError::Credentials(msg) => anyhow::anyhow!("Credentials error: {}", msg),
                 AuthError::TokenExchange(msg) => anyhow::anyhow!("Token exchange error: {}", msg),
             })?;
@@ -136,4 +142,22 @@ mod tests {
             "https://my-resource.openai.azure.com/openai"
         ));
     }
+
+    #[tokio::test]
+    async fn test_auth_header_bearer_token() {
+        let auth = AzureAuth::new(None, Some("my-token".to_string())).unwrap();
+        let provider = AzureAuthProvider { auth };
+        let (header, value) = provider.get_auth_header().await.unwrap();
+        assert_eq!(header, "Authorization");
+        assert_eq!(value, "Bearer my-token");
+    }
+
+    #[tokio::test]
+    async fn test_auth_header_api_key() {
+        let auth = AzureAuth::new(Some("my-key".to_string()), None).unwrap();
+        let provider = AzureAuthProvider { auth };
+        let (header, value) = provider.get_auth_header().await.unwrap();
+        assert_eq!(header, "api-key");
+        assert_eq!(value, "my-key");
+    }
 }

crates/goose/src/providers/azureauth.rs

@@ -31,6 +31,8 @@ pub struct AuthToken {
 pub enum AzureCredentials {
     /// API key based authentication
     ApiKey(String),
+    /// Pre-acquired Microsoft Entra ID access token (e.g. AZURE_OPENAI_AD_TOKEN)
+    BearerToken(String),
     /// Azure credential chain based authentication
     DefaultCredential,
 }
@@ -70,10 +72,11 @@ impl AzureAuth {
     ///
     /// # Returns
     /// * `Result<Self, AuthError>` - A new AzureAuth instance or an error if initialization fails
-    pub fn new(api_key: Option<String>) -> Result<Self, AuthError> {
-        let credentials = match api_key {
-            Some(key) => AzureCredentials::ApiKey(key),
-            None => AzureCredentials::DefaultCredential,
+    pub fn new(api_key: Option<String>, ad_token: Option<String>) -> Result<Self, AuthError> {
+        let credentials = match (ad_token, api_key) {
+            (Some(token), _) => AzureCredentials::BearerToken(token),
+            (None, Some(key)) => AzureCredentials::ApiKey(key),
+            (None, None) => AzureCredentials::DefaultCredential,
         };
 
         Ok(Self {
@@ -91,7 +94,8 @@ impl AzureAuth {
     ///
     /// This method implements an efficient token management strategy:
     /// 1. For API key auth, returns the API key directly
-    /// 2. For Azure credential chain:
+    /// 2. For bearer token auth, returns the pre-acquired token directly
+    /// 3. For Azure credential chain:
     ///    a. Checks the cache for a valid token
     ///    b. Returns the cached token if not expired
     ///    c. Obtains a new token if needed or expired
@@ -105,6 +109,10 @@ impl AzureAuth {
                 token_type: "Bearer".to_string(),
                 token_value: key.clone(),
             }),
+            AzureCredentials::BearerToken(token) => Ok(AuthToken {
+                token_type: "Bearer".to_string(),
+                token_value: token.clone(),
+            }),
             AzureCredentials::DefaultCredential => self.get_default_credential_token().await,
         }
     }
@@ -170,3 +178,43 @@ impl AzureAuth {
         Ok(auth_token)
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_ad_token_takes_precedence_over_api_key() {
+        let auth = AzureAuth::new(Some("key".to_string()), Some("token".to_string())).unwrap();
+        assert!(matches!(
+            auth.credential_type(),
+            AzureCredentials::BearerToken(_)
+        ));
+    }
+
+    #[test]
+    fn test_api_key_when_no_ad_token() {
+        let auth = AzureAuth::new(Some("key".to_string()), None).unwrap();
+        assert!(matches!(
+            auth.credential_type(),
+            AzureCredentials::ApiKey(_)
+        ));
+    }
+
+    #[test]
+    fn test_default_credential_when_neither() {
+        let auth = AzureAuth::new(None, None).unwrap();
+        assert!(matches!(
+            auth.credential_type(),
+            AzureCredentials::DefaultCredential
+        ));
+    }
+
+    #[tokio::test]
+    async fn test_bearer_token_get_token() {
+        let auth = AzureAuth::new(None, Some("my-token".to_string())).unwrap();
+        let token = auth.get_token().await.unwrap();
+        assert_eq!(token.token_type, "Bearer");
+        assert_eq!(token.token_value, "my-token");
+    }
+}