mirror of
https://github.com/aaif-goose/goose.git
synced 2026-07-03 14:10:03 +02:00
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:
@@ -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 {
|
||||
|
||||
@@ -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");
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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`
|
||||
|
||||
Reference in New Issue
Block a user