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>
This commit is contained in:
Hammad Khan
2026-06-17 07:09:05 +05:00
committed by GitHub
parent 9ca81e78d7
commit b3d3a8b29e
4 changed files with 94 additions and 13 deletions
@@ -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 {
+27 -3
View File
@@ -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");
}
}
+53 -5
View File
@@ -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");
}
}
@@ -27,7 +27,7 @@ goose is compatible with a wide range of LLM providers, allowing you to choose a
| [Anthropic](https://www.anthropic.com/) | Offers Claude, an advanced AI model for natural language tasks. | `ANTHROPIC_API_KEY`, `ANTHROPIC_HOST` (optional) |
| [Atomic Chat](https://github.com/AtomicBot-ai/Atomic-Chat) | Run local models with Atomic Chat's OpenAI-compatible server. **Because this provider runs locally, you must first [download a model](#local-llms).** | None required. Connects to local server at `localhost:1337` by default. |
| [Avian](https://avian.io/) | Cost-effective inference API with DeepSeek, Kimi, GLM, and MiniMax models. OpenAI-compatible with streaming and function calling support. | `AVIAN_API_KEY`, `AVIAN_HOST` (optional) |
| [Azure OpenAI](https://learn.microsoft.com/en-us/azure/ai-services/openai/) | Access Azure-hosted OpenAI models, including GPT-4 and GPT-3.5. Supports both API key and Azure credential chain authentication. | `AZURE_OPENAI_ENDPOINT`, `AZURE_OPENAI_DEPLOYMENT_NAME`, `AZURE_OPENAI_API_KEY` (optional) |
| [Azure OpenAI](https://learn.microsoft.com/en-us/azure/ai-services/openai/) | Access Azure-hosted OpenAI models, including GPT-4 and GPT-3.5. Supports API key, Entra ID bearer token, and Azure credential chain authentication. | `AZURE_OPENAI_ENDPOINT`, `AZURE_OPENAI_DEPLOYMENT_NAME`, `AZURE_OPENAI_API_KEY` (optional), `AZURE_OPENAI_AD_TOKEN` (optional) |
| [ChatGPT Codex](https://chatgpt.com/codex) | Access GPT-5 Codex models optimized for code generation and understanding. **Requires a ChatGPT Plus/Pro subscription.** | No manual key. Uses browser-based OAuth authentication for both CLI and Desktop. |
| [Databricks](https://www.databricks.com/) | Unified data analytics and AI platform for building and deploying models. | `DATABRICKS_HOST`, `DATABRICKS_TOKEN` |
| [Docker Model Runner](https://docs.docker.com/ai/model-runner/) | Local models running in Docker Desktop or Docker CE with OpenAI-compatible API endpoints. **Because this provider runs locally, you must first [download a model](#local-llms).** | `OPENAI_HOST`, `OPENAI_BASE_PATH` |
@@ -1363,12 +1363,15 @@ GitHub Copilot uses a device flow for authentication, so no API keys are require
4. Paste the code to authorize the application
5. When you return to goose, GitHub Copilot will be available as a provider in both CLI and Desktop.
## Azure OpenAI Credential Chain
## Azure OpenAI Authentication
goose supports two authentication methods for Azure OpenAI:
goose supports three authentication methods for Azure OpenAI:
1. **API Key Authentication** - Uses the `AZURE_OPENAI_API_KEY` for direct authentication
2. **Azure Credential Chain** - Uses Azure CLI credentials automatically without requiring an API key
1. **Entra ID Bearer Token** - Uses a pre-acquired Microsoft Entra access token from `AZURE_OPENAI_AD_TOKEN`, sent as `Authorization: Bearer <token>`. goose skips Azure CLI and token acquisition entirely, which suits enterprise deployments where only short-lived tokens are exposed to the runtime (e.g. obtained via `az account get-access-token --resource https://cognitiveservices.azure.com --query accessToken --output tsv`)
2. **API Key Authentication** - Uses the `AZURE_OPENAI_API_KEY` for direct authentication
3. **Azure Credential Chain** - Uses Azure CLI credentials automatically without requiring an API key
When more than one is configured, `AZURE_OPENAI_AD_TOKEN` takes precedence over `AZURE_OPENAI_API_KEY`, which takes precedence over the credential chain.
To use the Azure Credential Chain:
- Ensure you're logged in with `az login`