From 57d76ea171181fbaf93e793fb43663321924f01e Mon Sep 17 00:00:00 2001 From: Ming Wen Date: Sun, 24 May 2026 20:12:06 +0800 Subject: [PATCH 1/2] feat(azure-openai): wire AAD (Entra ID) client_credentials Bearer auth (#302 Phase F D6.6) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds the second Azure auth scheme to aisix-provider-azure-openai. Today the bridge supports only the resource-key scheme (`api-key:` header). This PR adds AAD client_credentials (Entra ID) so an operator can configure a ProviderKey backed by a service-principal app registration instead of pasting the resource's master api-key. Backward-compatible: existing api-key deployments keep working unchanged. The auth scheme is autodetected from the secret shape: - Secret starts with `{` → JSON-parse as AAD credentials {tenant_id, client_id, client_secret}. Bridge mints a token via the client_credentials grant, caches it, and sends Authorization: Bearer . - Otherwise → verbatim string, used as the resource api-key (sent via the api-key: header per the existing path). ## Wire shape (AAD branch) ``` POST https://login.microsoftonline.com/{tenant_id}/oauth2/v2.0/token Content-Type: application/x-www-form-urlencoded grant_type=client_credentials &client_id= &client_secret= &scope=https://cognitiveservices.azure.com/.default ``` Unlike Vertex SA OAuth (#387), AAD client_credentials is a straight form-encoded POST — NO JWT signing on the gateway side. No `jsonwebtoken` dep added; pure reqwest + serde. ## Cache Keyed by `(tenant_id, client_id)` — multiple ProviderKeys backed by the same AAD app share a slot, but distinct apps under the same tenant don't collide. Refresh 60s before upstream-reported expiry. ## Error classification (audit-aware) Mirrors the Vertex audit MEDIUM fix from ai-gateway#387: - AAD 5xx → BridgeError::UpstreamStatus + Retry-After propagated (transient backend should hit cooldown layer, not 500 operator-must-fix). - AAD 4xx → BridgeError::Config (invalid_client / revoked secret / wrong scope IS operator-actionable). ## Files - `aad_token_mint.rs` (new, ~290 lines + tests): TokenMinter, AadCredentials, RwLock-backed cache. 7 unit tests covering happy mint, cache reuse, cache separation across distinct apps, 5xx/4xx classification, empty/URL-injection rejection at validate(). - `bridge.rs`: - Added `AzureSecret` discriminated parse (api-key verbatim vs AAD JSON), and `AzureAuth` resolved-header pair. - Bridge struct carries an Arc for the AAD path. - `resolve_auth(ctx)` is called BEFORE the chat / chat_stream future so AAD mint failures surface as direct Err returns (matches existing 4xx/timeout error semantics). - `build_request_headers` signature changed from `&str` to `&AzureAuth`; emits either `api-key:` (legacy) or `Authorization: Bearer` (AAD) based on which is set. - Added test-only `with_aad_token_endpoint_override` seam mirroring the existing `with_url_override` pattern. - Removed the now-unused `fn api_key()` helper (replaced by `AzureSecret::parse`). - 7 new tests: secret-parse (api-key / AAD / empty / bad JSON), end-to-end chat with AAD bearer header set, cache reuse across 3 chats, AAD 4xx surfaces before Azure call. - `lib.rs`: declares aad_token_mint module, ticks D6.6 in status block. `cargo test -p aisix-provider-azure-openai` → 53/53 PASS (+7). `cargo clippy -p aisix-provider-azure-openai --all-targets -- -D warnings` clean. `cargo fmt --all` applied. ## References (CLAUDE.md §7) - Microsoft identity platform — client credentials grant flow: https://learn.microsoft.com/en-us/entra/identity-platform/v2-oauth2-client-creds-grant-flow - Azure OpenAI Entra ID auth: https://learn.microsoft.com/en-us/azure/ai-services/openai/how-to/managed-identity - OAuth2 RFC 6749 §4.4 (client_credentials): https://www.rfc-editor.org/rfc/rfc6749#section-4.4 - Mirrors the audit-corrected pattern from `aisix-provider-vertex::token_mint` (ai-gateway#387). ## Unblocks AC.12 hardening in api7/AISIX-Cloud#302: Azure-OpenAI was already ~70% done (chat + stream + filter tolerance via #319); the AAD auth path was the explicit D6.6 gap called out in the audit. With this PR Phase F is complete. Live e2e against the Step 0.1 mock-llm Azure profile is the next sub-step (separate PR). --- .../src/aad_token_mint.rs | 455 ++++++++++++++++++ .../aisix-provider-azure-openai/src/bridge.rs | 432 +++++++++++++++-- crates/aisix-provider-azure-openai/src/lib.rs | 16 +- 3 files changed, 862 insertions(+), 41 deletions(-) create mode 100644 crates/aisix-provider-azure-openai/src/aad_token_mint.rs diff --git a/crates/aisix-provider-azure-openai/src/aad_token_mint.rs b/crates/aisix-provider-azure-openai/src/aad_token_mint.rs new file mode 100644 index 00000000..1afb95e0 --- /dev/null +++ b/crates/aisix-provider-azure-openai/src/aad_token_mint.rs @@ -0,0 +1,455 @@ +//! Azure AD (Entra ID) `client_credentials` OAuth2 flow + token cache +//! for the Azure OpenAI Service bridge. +//! +//! Unlike Vertex (JWT-bearer assertion grant with RS256 signing), +//! Azure AAD's `client_credentials` flow is a straight form-encoded +//! POST — no JWT signing on the gateway side: +//! +//! ```text +//! POST https://login.microsoftonline.com/{tenant_id}/oauth2/v2.0/token +//! Content-Type: application/x-www-form-urlencoded +//! +//! grant_type=client_credentials +//! &client_id= +//! &client_secret= +//! &scope=https://cognitiveservices.azure.com/.default +//! ``` +//! +//! Response: `{access_token, expires_in, token_type: "Bearer"}`. The +//! customer then sends `Authorization: Bearer ` on every +//! Azure OpenAI request — replacing the `api-key:` header used by the +//! resource-key scheme. +//! +//! Cache is keyed by `(tenant_id, client_id)` because two distinct +//! AAD app registrations under the same tenant must NOT share a +//! token slot. Refresh ~60s before upstream-reported expiry so an +//! in-flight request never lands on an expired token. +//! +//! 5xx from the token endpoint surfaces as `BridgeError::UpstreamStatus` +//! with `Retry-After` propagated (transient AAD outage should hit the +//! cooldown layer, not the operator-must-fix path). 4xx remains +//! `BridgeError::Config` — `invalid_client` / `unauthorized_client` / +//! revoked secret IS operator-actionable. Pattern mirrors +//! `aisix-provider-vertex::token_mint` after the audit on +//! ai-gateway#387. +//! +//! # References +//! +//! - Microsoft identity platform — client credentials grant flow: +//! +//! - Azure OpenAI authentication (Entra ID section): +//! +//! - OAuth2 `client_credentials` spec (RFC 6749 §4.4): +//! + +use aisix_gateway::BridgeError; +use reqwest::Client; +use serde::Deserialize; +use std::collections::HashMap; +use std::sync::Arc; +use std::time::{Duration, Instant}; +use tokio::sync::RwLock; + +/// Default scope for Azure OpenAI Service. Per Microsoft's docs the +/// `.default` suffix tells AAD to issue a token covering every +/// permission the app registration is configured for — the standard +/// pattern for non-interactive service-to-service auth. +const AZURE_OPENAI_SCOPE: &str = "https://cognitiveservices.azure.com/.default"; + +/// Refresh cached tokens at least this many seconds before their +/// reported expiry. Prevents a request from picking up a token that +/// expires while the request is mid-flight. +const TOKEN_REFRESH_SAFETY_MARGIN: Duration = Duration::from_secs(60); + +/// AAD app-registration credentials for the `client_credentials` +/// grant. Field names match the JSON shape an operator pastes +/// directly from the Azure portal's "Certificates & secrets" view. +#[derive(Debug, Clone, Deserialize)] +pub(crate) struct AadCredentials { + /// AAD tenant UUID (or vanity domain). Embedded in the token + /// endpoint URL path. + pub tenant_id: String, + /// App-registration (application) UUID. + pub client_id: String, + /// Client secret value (NOT the secret id). Confidential. + pub client_secret: String, +} + +impl AadCredentials { + /// Cheap shape validation at parse time. Tenant id / client id + /// must not contain URL-control characters (so we can interpolate + /// the tenant into the token endpoint path without injection). + pub fn validate(&self) -> Result<(), BridgeError> { + for (name, value) in [ + ("tenant_id", &self.tenant_id), + ("client_id", &self.client_id), + ("client_secret", &self.client_secret), + ] { + if value.is_empty() { + return Err(BridgeError::Config(format!( + "azure aad credentials.{name} is empty" + ))); + } + } + for (name, value) in [ + ("tenant_id", &self.tenant_id), + ("client_id", &self.client_id), + ] { + if value.contains('/') + || value.contains('?') + || value.contains('#') + || value.contains(' ') + || value.contains('\t') + || value.contains('\n') + || value.contains("..") + { + return Err(BridgeError::Config(format!( + "azure aad credentials.{name} {value:?} contains URL-control \ + characters — reject `/`, `?`, `#`, whitespace, `..`" + ))); + } + } + Ok(()) + } +} + +/// Token-endpoint response shape per OAuth2 spec. +#[derive(Deserialize)] +struct TokenResponse { + access_token: String, + expires_in: u64, +} + +#[derive(Clone)] +struct CachedToken { + access_token: String, + expires_at: Instant, +} + +/// In-process token cache + minter. One instance per +/// `AzureOpenAiBridge`. Cache is keyed by `(tenant_id, client_id)` +/// so multiple ProviderKeys backed by the same AAD app share a slot, +/// but distinct apps under the same tenant don't collide. +pub(crate) struct TokenMinter { + client: Client, + cache: Arc>>, + /// Test-only override for the AAD token endpoint. In production + /// the URL is derived from the tenant id; tests substitute a + /// wiremock URI here. + #[cfg(test)] + token_endpoint_override: Option, +} + +impl TokenMinter { + pub fn new(client: Client) -> Self { + Self { + client, + cache: Arc::new(RwLock::new(HashMap::new())), + #[cfg(test)] + token_endpoint_override: None, + } + } + + /// Test-only seam: replace the `login.microsoftonline.com` host + /// with this URL. Tenant id is still interpolated into the path + /// (so the request URL shape is verifiable end-to-end against + /// wiremock matchers). + #[cfg(test)] + pub(crate) fn with_token_endpoint_override(mut self, url: impl Into) -> Self { + self.token_endpoint_override = Some(url.into()); + self + } + + /// Resolve an access token for `creds`. Returns the cached token + /// when one exists and is unexpired; otherwise mints a fresh one + /// and caches it. + pub async fn get_token(&self, creds: &AadCredentials) -> Result { + let key = (creds.tenant_id.clone(), creds.client_id.clone()); + { + let cache = self.cache.read().await; + if let Some(cached) = cache.get(&key) { + if cached.expires_at > Instant::now() { + return Ok(cached.access_token.clone()); + } + } + } + let (access_token, expires_in_secs) = self.mint(creds).await?; + let cached = CachedToken { + access_token: access_token.clone(), + expires_at: Instant::now() + + Duration::from_secs(expires_in_secs).saturating_sub(TOKEN_REFRESH_SAFETY_MARGIN), + }; + self.cache.write().await.insert(key, cached); + Ok(access_token) + } + + /// Mint a fresh token by POSTing the client_credentials grant + /// to the AAD token endpoint. Returns `(access_token, expires_in)`. + async fn mint(&self, creds: &AadCredentials) -> Result<(String, u64), BridgeError> { + creds.validate()?; + let endpoint = self.resolve_token_endpoint(creds); + let resp = self + .client + .post(&endpoint) + .form(&[ + ("grant_type", "client_credentials"), + ("client_id", &creds.client_id), + ("client_secret", &creds.client_secret), + ("scope", AZURE_OPENAI_SCOPE), + ]) + .send() + .await + .map_err(|e| { + BridgeError::Transport(format!("azure aad token mint POST {endpoint}: {e}")) + })?; + + let status = resp.status(); + if !status.is_success() { + // Read Retry-After BEFORE consuming the body. + let retry_after = aisix_gateway::parse_retry_after(resp.headers()); + let body = resp.text().await.unwrap_or_default(); + // Cap body to 500 chars. AAD error envelopes + // `{error: "invalid_client", error_description: "..."}` are + // operator-actionable; cap is defense-in-depth against a + // mis-deployed front-door that returns an HTML error page. + let truncated: String = body.chars().take(500).collect(); + let msg = format!("azure aad token mint upstream returned HTTP {status}: {truncated}"); + // Mirror Vertex audit MEDIUM (ai-gateway#387): 5xx is + // transient upstream — should surface UpstreamStatus + // (cooldown semantics) with the Retry-After hint. 4xx + // (invalid_client / unauthorized_client / revoked secret) + // is operator-actionable, stays Config. + return Err(if status.is_server_error() { + BridgeError::upstream_status_with_retry_after(status.as_u16(), msg, retry_after) + } else { + BridgeError::Config(msg) + }); + } + let parsed: TokenResponse = resp.json().await.map_err(|e| { + BridgeError::UpstreamDecode(format!("azure aad token mint response: {e}")) + })?; + Ok((parsed.access_token, parsed.expires_in)) + } + + fn resolve_token_endpoint(&self, creds: &AadCredentials) -> String { + #[cfg(test)] + if let Some(base) = &self.token_endpoint_override { + // Tests get a fixed URL — tenant id is still validated + // upstream of this call. + return base.clone(); + } + format!( + "https://login.microsoftonline.com/{}/oauth2/v2.0/token", + creds.tenant_id + ) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use wiremock::matchers::{body_string_contains, method}; + use wiremock::{Mock, MockServer, Request, ResponseTemplate}; + + fn sample_creds() -> AadCredentials { + AadCredentials { + tenant_id: "11111111-1111-1111-1111-111111111111".into(), + client_id: "22222222-2222-2222-2222-222222222222".into(), + client_secret: "fake-secret-not-a-real-one".into(), + } + } + + #[tokio::test] + async fn mint_posts_client_credentials_grant_with_correct_scope() { + let server = MockServer::start().await; + let captured_body: Arc>> = + Arc::new(std::sync::Mutex::new(None)); + let captured_for_responder = captured_body.clone(); + Mock::given(method("POST")) + .and(body_string_contains("grant_type=client_credentials")) + .respond_with(move |req: &Request| { + *captured_for_responder.lock().unwrap() = + Some(String::from_utf8(req.body.clone()).unwrap()); + ResponseTemplate::new(200).set_body_json(serde_json::json!({ + "access_token": "aad.minted-by-mock", + "expires_in": 3600, + "token_type": "Bearer" + })) + }) + .mount(&server) + .await; + + let minter = TokenMinter::new(Client::new()).with_token_endpoint_override(server.uri()); + let token = minter.get_token(&sample_creds()).await.unwrap(); + assert_eq!(token, "aad.minted-by-mock"); + + let body = captured_body + .lock() + .unwrap() + .clone() + .expect("body captured"); + // Form fields per RFC 6749 §4.4 + Microsoft client-credentials docs. + assert!(body.contains("grant_type=client_credentials")); + assert!(body.contains("client_id=22222222-2222-2222-2222-222222222222")); + assert!(body.contains("client_secret=fake-secret-not-a-real-one")); + assert!(body.contains("scope=https%3A%2F%2Fcognitiveservices.azure.com%2F.default")); + } + + #[tokio::test] + async fn get_token_caches_within_ttl_and_only_calls_endpoint_once() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ + "access_token": "aad.cache-hit", + "expires_in": 3600, + "token_type": "Bearer" + }))) + .expect(1) + .mount(&server) + .await; + let minter = TokenMinter::new(Client::new()).with_token_endpoint_override(server.uri()); + let creds = sample_creds(); + for _ in 0..3 { + assert_eq!(minter.get_token(&creds).await.unwrap(), "aad.cache-hit"); + } + } + + #[tokio::test] + async fn cache_separates_distinct_app_registrations_under_same_tenant() { + // Two apps under the same tenant get distinct cache slots — + // a regression that keyed only by tenant_id would serve + // app B's request app A's token (or vice versa). Pin the + // separation explicitly. + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(body_string_contains("client_id=aaaa")) + .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ + "access_token": "token-for-app-A", + "expires_in": 3600, + "token_type": "Bearer" + }))) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(body_string_contains("client_id=bbbb")) + .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ + "access_token": "token-for-app-B", + "expires_in": 3600, + "token_type": "Bearer" + }))) + .mount(&server) + .await; + + let minter = TokenMinter::new(Client::new()).with_token_endpoint_override(server.uri()); + let app_a = AadCredentials { + tenant_id: "shared-tenant".into(), + client_id: "aaaa-aaaa".into(), + client_secret: "secret-a".into(), + }; + let app_b = AadCredentials { + tenant_id: "shared-tenant".into(), + client_id: "bbbb-bbbb".into(), + client_secret: "secret-b".into(), + }; + assert_eq!(minter.get_token(&app_a).await.unwrap(), "token-for-app-A"); + assert_eq!(minter.get_token(&app_b).await.unwrap(), "token-for-app-B"); + // Repeat — both must come from cache, not re-mint. + assert_eq!(minter.get_token(&app_a).await.unwrap(), "token-for-app-A"); + assert_eq!(minter.get_token(&app_b).await.unwrap(), "token-for-app-B"); + } + + #[tokio::test] + async fn token_endpoint_5xx_surfaces_as_upstream_status_with_retry_hint() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .respond_with( + ResponseTemplate::new(503) + .insert_header("Retry-After", "45") + .set_body_string("Service Unavailable: please retry"), + ) + .mount(&server) + .await; + let minter = TokenMinter::new(Client::new()).with_token_endpoint_override(server.uri()); + let err = minter.get_token(&sample_creds()).await.err().unwrap(); + match err { + BridgeError::UpstreamStatus { + status, + retry_after, + .. + } => { + assert_eq!(status, 503); + assert_eq!(retry_after, Some(Duration::from_secs(45))); + } + other => panic!("expected UpstreamStatus, got {other:?}"), + } + } + + #[tokio::test] + async fn token_endpoint_4xx_invalid_client_surfaces_as_config_error() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .respond_with(ResponseTemplate::new(401).set_body_string( + r#"{"error":"invalid_client","error_description":"AADSTS7000215: Invalid client secret provided"}"#, + )) + .mount(&server) + .await; + let minter = TokenMinter::new(Client::new()).with_token_endpoint_override(server.uri()); + let err = minter.get_token(&sample_creds()).await.err().unwrap(); + match err { + BridgeError::Config(msg) => { + assert!(msg.contains("HTTP 401")); + assert!(msg.contains("invalid_client")); + } + other => panic!("expected Config error, got {other:?}"), + } + } + + #[tokio::test] + async fn empty_tenant_id_rejected_before_endpoint_call() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .respond_with(ResponseTemplate::new(200)) + .expect(0) + .mount(&server) + .await; + let minter = TokenMinter::new(Client::new()).with_token_endpoint_override(server.uri()); + let creds = AadCredentials { + tenant_id: "".into(), + client_id: "abc".into(), + client_secret: "xyz".into(), + }; + let err = minter.get_token(&creds).await.err().unwrap(); + match err { + BridgeError::Config(msg) => { + assert!(msg.contains("tenant_id is empty")); + } + other => panic!("expected Config, got {other:?}"), + } + } + + #[tokio::test] + async fn tenant_id_with_url_injection_rejected_before_endpoint_call() { + // A `/` in tenant_id could redirect the token POST to a path + // the operator never authorized. validate() must catch this + // before we interpolate into the URL. + let server = MockServer::start().await; + Mock::given(method("POST")) + .respond_with(ResponseTemplate::new(200)) + .expect(0) + .mount(&server) + .await; + let minter = TokenMinter::new(Client::new()).with_token_endpoint_override(server.uri()); + let creds = AadCredentials { + tenant_id: "../malicious".into(), + client_id: "abc".into(), + client_secret: "xyz".into(), + }; + let err = minter.get_token(&creds).await.err().unwrap(); + match err { + BridgeError::Config(msg) => { + assert!(msg.contains("URL-control")); + } + other => panic!("expected Config, got {other:?}"), + } + } +} diff --git a/crates/aisix-provider-azure-openai/src/bridge.rs b/crates/aisix-provider-azure-openai/src/bridge.rs index 36caff52..5f887ef0 100644 --- a/crates/aisix-provider-azure-openai/src/bridge.rs +++ b/crates/aisix-provider-azure-openai/src/bridge.rs @@ -39,8 +39,11 @@ use aisix_provider_openai::wire::{ OpenAiResponse, OpenAiStreamChunk, }; +use crate::aad_token_mint::TokenMinter; use crate::wire; +use std::sync::Arc; + /// Family Bridge for Azure OpenAI Service. pub struct AzureOpenAiBridge { client: Client, @@ -48,6 +51,12 @@ pub struct AzureOpenAiBridge { /// so dashboards can split Azure traffic from canonical OpenAI /// traffic in metrics. name: &'static str, + /// In-process AAD token cache + minter. Used only when the + /// inbound `ProviderKey.secret` parses to the AAD branch; the + /// api-key branch bypasses this entirely. `Arc` so the bridge + /// remains cheaply clonable for callers that share it across + /// Hub registrations. + token_minter: Arc, /// Test-only POST URL override. When set, [`Bridge::chat`] / /// [`Bridge::chat_stream`] still run resolve / validation / /// header / body building against the real `AzureUpstreamRef`, @@ -72,6 +81,7 @@ impl AzureOpenAiBridge { /// timeouts. Public surface — not test-only. pub fn with_client(client: Client) -> Self { Self { + token_minter: Arc::new(TokenMinter::new(client.clone())), client, name: "azure-openai", #[cfg(test)] @@ -79,6 +89,17 @@ impl AzureOpenAiBridge { } } + /// Test-only seam: replace the AAD token endpoint host on the + /// internal minter (without touching the chat-completions URL + /// override on `url_override`). Used by AAD-flow tests so + /// `client_credentials` POSTs land on a wiremock instance. + #[cfg(test)] + pub(crate) fn with_aad_token_endpoint_override(mut self, url: impl Into) -> Self { + self.token_minter = + Arc::new(TokenMinter::new(self.client.clone()).with_token_endpoint_override(url)); + self + } + /// Resolve the URL the bridge will POST to. Returns /// `upstream.chat_completions_url()` in production; tests can /// override via [`Self::with_url_override`]. @@ -99,6 +120,28 @@ impl AzureOpenAiBridge { self.url_override = Some(url.into()); self } + + /// Resolve the per-request auth pair from the provider key's + /// secret. Returns either an `api_key` (verbatim resource-key) + /// or a `bearer_token` (freshly minted-or-cached AAD access + /// token). Token-mint failures (e.g. AAD 401 invalid_client) + /// surface as `Err` here so `chat()` / `chat_stream()` short- + /// circuit BEFORE attempting the upstream Azure OpenAI call. + async fn resolve_auth(&self, ctx: &BridgeContext) -> Result { + match AzureSecret::parse(&ctx.provider_key.secret)? { + AzureSecret::ApiKey(key) => Ok(AzureAuth { + api_key: Some(key), + bearer_token: None, + }), + AzureSecret::Aad(creds) => { + let token = self.token_minter.get_token(&creds).await?; + Ok(AzureAuth { + api_key: None, + bearer_token: Some(token), + }) + } + } + } } impl Default for AzureOpenAiBridge { @@ -243,16 +286,61 @@ fn validate_url_token(name: &str, value: &str) -> Result<(), BridgeError> { Ok(()) } -/// Pull the api-key from the BridgeContext's ProviderKey. -fn api_key(ctx: &BridgeContext) -> Result<&str, BridgeError> { - let k = &ctx.provider_key.secret; - if k.is_empty() { - Err(BridgeError::Config("provider_key.secret is empty".into())) - } else { - Ok(k.as_str()) +/// Discriminated auth scheme. Today's resource-key deployments keep +/// the verbatim-string secret shape (backward-compat); the AAD path +/// is opted into by encoding the secret as JSON +/// `{tenant_id, client_id, client_secret}`. +/// +/// Detection is by the leading character of the trimmed secret — +/// `{` triggers JSON parse, anything else is treated as a literal +/// api-key. Real Azure api-keys are URL-safe base64 strings, never +/// starting with `{`, so the heuristic is unambiguous in practice. +#[derive(Debug)] +pub(crate) enum AzureSecret { + ApiKey(String), + Aad(crate::aad_token_mint::AadCredentials), +} + +impl AzureSecret { + /// Parse the inbound `ProviderKey.secret` into either the api-key + /// or AAD branch. + /// + /// **Audit-aware:** error messages MUST NOT echo raw secret + /// bytes. The branch heuristic is `starts_with('{')` after + /// trimming — that test does not panic, returns Bool, no leaks. + /// The JSON parse error message is fixed (not interpolated from + /// serde) so partial secret contents can't surface in the error. + pub(crate) fn parse(secret: &str) -> Result { + let trimmed = secret.trim(); + if trimmed.is_empty() { + return Err(BridgeError::Config("provider_key.secret is empty".into())); + } + if trimmed.starts_with('{') { + let creds: crate::aad_token_mint::AadCredentials = serde_json::from_str(trimmed) + .map_err(|_e| { + BridgeError::Config( + "azure provider_key.secret looks JSON-shaped but failed to parse \ + as AAD client_credentials \ + {tenant_id, client_id, client_secret}" + .into(), + ) + })?; + creds.validate()?; + Ok(AzureSecret::Aad(creds)) + } else { + Ok(AzureSecret::ApiKey(trimmed.to_string())) + } } } +/// Resolved per-request auth header pair the bridge writes onto the +/// outbound request. Exactly ONE of `api_key` / `bearer_token` is +/// `Some`; the other is `None`. +pub(crate) struct AzureAuth { + pub api_key: Option, + pub bearer_token: Option, +} + /// Pull the upstream deployment name off the BridgeContext. Azure /// deployment names (operator-defined in the Azure portal, e.g. /// `gpt4o-prod`) live on Model.model_name. `req.model` is the @@ -401,32 +489,61 @@ fn prepare_outbound_body( Ok(body) } -/// Build the base outbound `HeaderMap` for Azure: -/// - `api-key: ` (Azure's standard auth header — NOT -/// `Authorization: Bearer`) -/// - `Content-Type: application/json` -/// - `x-aisix-request-id: ` -/// - `Accept: text/event-stream` when streaming +/// Build the base outbound `HeaderMap` for Azure. The auth header +/// depends on the resolved auth scheme: /// -/// Bridge-owned headers are inserted before `apply_default_headers` so -/// the reserved-headers list in `aisix-provider-openai::overrides` -/// (which already covers `api-key`, `authorization`, `x-api-key`, plus -/// hop-by-hop / proxy-auth headers) cannot overwrite them. Defense in -/// depth: the reserved-list blocks even before the -/// `headers.contains_key` guard inside `apply_default_headers`. +/// - api-key scheme → `api-key: ` (Azure docs: +/// ; +/// the literal lowercase-hyphenated `api-key`, NOT +/// `Authorization: Bearer`). +/// - AAD scheme → `Authorization: Bearer ` (the +/// industry-standard Bearer header; matches how Azure's own +/// Python SDK sets the header on Entra ID auth). +/// +/// In both branches the bridge also sets `Content-Type: application/json`, +/// `x-aisix-request-id: `, and (for streaming) +/// `Accept: text/event-stream`. Bridge-owned headers are inserted +/// before `apply_default_headers` so the reserved-headers list in +/// `aisix-provider-openai::overrides` (which already covers +/// `api-key`, `authorization`, `x-api-key`, plus hop-by-hop / +/// proxy-auth headers) cannot overwrite them. Defense in depth: the +/// reserved-list blocks even before the `headers.contains_key` guard +/// inside `apply_default_headers`. fn build_request_headers( - api_key_str: &str, + auth: &AzureAuth, request_id: &str, sse: bool, request: Option<&RequestOverrides>, ) -> Result { let mut headers = HeaderMap::new(); - let api_key_value = HeaderValue::from_str(api_key_str) - .map_err(|e| BridgeError::Config(format!("api key contains invalid header chars: {e}")))?; - // Per Azure docs (https://learn.microsoft.com/en-us/azure/ai-services/openai/reference) - // the canonical auth header for the api-key scheme is the literal - // `api-key` (lowercase, hyphenated). - headers.insert(HeaderName::from_static("api-key"), api_key_value); + match (&auth.api_key, &auth.bearer_token) { + (Some(key), None) => { + let value = HeaderValue::from_str(key).map_err(|e| { + BridgeError::Config(format!("api key contains invalid header chars: {e}")) + })?; + headers.insert(HeaderName::from_static("api-key"), value); + } + (None, Some(token)) => { + // Bearer token from AAD. Same byte-validation pattern as + // api-key. Note: the validation error MUST NOT echo the + // token bytes — `http::InvalidHeaderValue`'s Display is + // opaque, but a future change could include byte + // position; rebind to a fixed message to be safe. + let value = HeaderValue::from_str(&format!("Bearer {token}")).map_err(|_| { + BridgeError::Config( + "azure aad access_token contains invalid header characters".into(), + ) + })?; + headers.insert(header::AUTHORIZATION, value); + } + // parse() / resolve_auth() ensure exactly one is Some — keep + // explicit guard for defense in depth. + _ => { + return Err(BridgeError::Config( + "internal: AzureAuth must set exactly one of api_key / bearer_token".into(), + )) + } + } headers.insert( header::CONTENT_TYPE, HeaderValue::from_static("application/json"), @@ -466,7 +583,10 @@ impl Bridge for AzureOpenAiBridge { let _ = wire::reserved_query_params(); let _ = wire::reserved_auth_headers(); - let key = api_key(ctx)?; + // Resolve auth BEFORE entering the request future so AAD + // mint failures surface as a direct Err return rather than + // a transport timeout. api-key path is a no-op string copy. + let auth = self.resolve_auth(ctx).await?; // Azure expects the deployment name in the URL path; the JSON // body's `model` field is ignored by Azure (or echoed back). // We still set it to the deployment name for log-trace clarity @@ -479,7 +599,7 @@ impl Bridge for AzureOpenAiBridge { ctx.provider_key.response.as_ref(), )?; let headers = build_request_headers( - key, + &auth, &ctx.request_id, false, ctx.provider_key.request.as_ref(), @@ -526,7 +646,8 @@ impl Bridge for AzureOpenAiBridge { let _ = wire::reserved_query_params(); let _ = wire::reserved_auth_headers(); - let key = api_key(ctx)?; + // See chat() — resolve auth before the request future. + let auth = self.resolve_auth(ctx).await?; let messages = messages_from(req); let typed = build_request(req, deployment, &messages, true); let body = prepare_outbound_body( @@ -535,7 +656,7 @@ impl Bridge for AzureOpenAiBridge { ctx.provider_key.response.as_ref(), )?; let headers = build_request_headers( - key, + &auth, &ctx.request_id, true, ctx.provider_key.request.as_ref(), @@ -883,6 +1004,20 @@ mod tests { ) } + /// Build a `ProviderKey` whose `secret` is the JSON-encoded AAD + /// credentials shape. Used by the D6.6 (Entra ID) tests; sidesteps + /// the string-escaping awkwardness of embedding JSON inside JSON. + fn sample_pk_with_aad_secret(api_base: &str) -> Arc { + let aad_json = r#"{"tenant_id":"tenant-uuid-aaa","client_id":"client-uuid-bbb","client_secret":"aad-secret-rotation-managed"}"#; + let pk = serde_json::from_value::(serde_json::json!({ + "display_name": "azure-aad-prod", + "secret": aad_json, + "api_base": api_base, + })) + .unwrap(); + Arc::new(pk) + } + fn sample_pk_with_overrides(api_base: &str, overrides_json: &str) -> Arc { Arc::new( serde_json::from_str(&format!( @@ -912,6 +1047,16 @@ mod tests { /// what `AzureUpstreamRef::chat_completions_url()` would produce /// but rooted at the mock's URI. Pass to /// [`AzureOpenAiBridge::with_url_override`]. + /// Test helper: build an [`AzureAuth`] for the api-key branch. + /// Mirrors the inbound shape the production parser produces for + /// a verbatim-string secret. + fn api_key_auth(key: &str) -> AzureAuth { + AzureAuth { + api_key: Some(key.to_string()), + bearer_token: None, + } + } + fn mock_chat_url(mock_uri: &str, deployment: &str) -> String { format!( "{}/openai/deployments/{}/chat/completions?api-version=2024-10-21", @@ -923,7 +1068,8 @@ mod tests { fn build_request_headers_uses_api_key_not_bearer() { // Critical Azure-vs-OpenAI distinction: the auth header is // literally `api-key`, NOT `Authorization: Bearer`. - let headers = build_request_headers("az-secret-key", "req-1", false, None).unwrap(); + let headers = + build_request_headers(&api_key_auth("az-secret-key"), "req-1", false, None).unwrap(); assert_eq!(headers.get("api-key").unwrap(), "az-secret-key"); assert!( !headers.contains_key("authorization"), @@ -939,7 +1085,7 @@ mod tests { #[test] fn build_request_headers_sets_sse_accept_when_streaming() { - let headers = build_request_headers("az-key", "req-1", true, None).unwrap(); + let headers = build_request_headers(&api_key_auth("az-key"), "req-1", true, None).unwrap(); assert_eq!(headers.get("accept").unwrap(), "text/event-stream"); } @@ -960,8 +1106,13 @@ mod tests { default_body_fields: Default::default(), default_headers, }; - let headers = - build_request_headers("legit-key", "req-1", false, Some(&request_overrides)).unwrap(); + let headers = build_request_headers( + &api_key_auth("legit-key"), + "req-1", + false, + Some(&request_overrides), + ) + .unwrap(); assert_eq!( headers.get("api-key").unwrap(), "legit-key", @@ -984,7 +1135,9 @@ mod tests { default_body_fields: Default::default(), default_headers, }; - let headers = build_request_headers("k", "req-1", false, Some(&request_overrides)).unwrap(); + let headers = + build_request_headers(&api_key_auth("k"), "req-1", false, Some(&request_overrides)) + .unwrap(); assert_eq!(headers.get("x-custom-trace").unwrap(), "trace-123"); } @@ -992,13 +1145,15 @@ mod tests { fn build_request_headers_rejects_invalid_api_key_chars() { // A secret with a newline would let an operator inject extra // headers via the api-key value. - let err = build_request_headers("legit\nx-evil: 1", "req-1", false, None).unwrap_err(); + let err = build_request_headers(&api_key_auth("legit\nx-evil: 1"), "req-1", false, None) + .unwrap_err(); assert!(matches!(err, BridgeError::Config(_))); } #[test] fn build_request_headers_rejects_invalid_request_id_chars() { - let err = build_request_headers("legit", "req\nbad", false, None).unwrap_err(); + let err = + build_request_headers(&api_key_auth("legit"), "req\nbad", false, None).unwrap_err(); assert!(matches!(err, BridgeError::Config(_))); } @@ -1663,4 +1818,207 @@ mod tests { Err(other) => panic!("expected Timeout, got {other:?}"), } } + + // ─── AAD (Entra ID) auth scheme (D6.6) ───────────────────────── + + /// JSON-shaped Azure secret triggering the AAD branch. + fn aad_secret_json() -> String { + r#"{ + "tenant_id": "tenant-uuid-aaa", + "client_id": "client-uuid-bbb", + "client_secret": "aad-secret-rotation-managed" + }"# + .to_string() + } + + #[test] + fn azure_secret_parses_verbatim_string_as_api_key() { + // Backward compat: existing pre-D6.6 deployments encode the + // resource api-key as a bare string. + let parsed = AzureSecret::parse("legacy-api-key-string").unwrap(); + match parsed { + AzureSecret::ApiKey(k) => assert_eq!(k, "legacy-api-key-string"), + AzureSecret::Aad(_) => panic!("verbatim string must parse as api-key, not AAD"), + } + } + + #[test] + fn azure_secret_parses_json_object_as_aad_credentials() { + let parsed = AzureSecret::parse(&aad_secret_json()).unwrap(); + match parsed { + AzureSecret::Aad(creds) => { + assert_eq!(creds.tenant_id, "tenant-uuid-aaa"); + assert_eq!(creds.client_id, "client-uuid-bbb"); + assert_eq!(creds.client_secret, "aad-secret-rotation-managed"); + } + AzureSecret::ApiKey(_) => panic!("JSON-shaped secret must parse as AAD, not api-key"), + } + } + + #[test] + fn azure_secret_rejects_empty_secret() { + let err = AzureSecret::parse(" ").unwrap_err(); + assert!(matches!(err, BridgeError::Config(_))); + } + + #[test] + fn azure_secret_rejects_json_missing_required_aad_fields() { + let err = AzureSecret::parse(r#"{"tenant_id":"t"}"#).unwrap_err(); + match err { + BridgeError::Config(msg) => { + // Message must NOT echo raw secret bytes (audit-aware). + assert!(msg.contains("looks JSON-shaped")); + assert!(!msg.contains("tenant-uuid-aaa")); + } + other => panic!("expected Config, got {other:?}"), + } + } + + #[tokio::test] + async fn chat_with_aad_secret_mints_token_and_sets_authorization_bearer_header() { + // End-to-end pin: AAD JSON secret → token mint via wiremock + // AAD endpoint → chat POST to Azure OpenAI endpoint carrying + // Authorization: Bearer (NOT api-key:). + let aad_server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/oauth2/v2.0/token")) + .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ + "access_token": "aad.bearer.test-token", + "expires_in": 3600, + "token_type": "Bearer" + }))) + .expect(1) + .mount(&aad_server) + .await; + + let azure_server = MockServer::start().await; + let captured_headers: std::sync::Arc>> = + std::sync::Arc::new(std::sync::Mutex::new(None)); + let captured_for_responder = captured_headers.clone(); + Mock::given(method("POST")) + .respond_with(move |req: &wiremock::Request| { + *captured_for_responder.lock().unwrap() = Some(req.headers.clone()); + ResponseTemplate::new(200).set_body_json(serde_json::json!({ + "id": "chatcmpl-aad-test", + "object": "chat.completion", + "created": 1700000000_i64, + "model": "gpt-4o", + "choices": [{ + "index": 0, + "message": {"role": "assistant", "content": "hello via aad"}, + "finish_reason": "stop" + }], + "usage": {"prompt_tokens": 4, "completion_tokens": 6, "total_tokens": 10} + })) + }) + .mount(&azure_server) + .await; + + let bridge = AzureOpenAiBridge::new() + .with_url_override(mock_chat_url(&azure_server.uri(), "gpt4o-prod")) + .with_aad_token_endpoint_override(format!("{}/oauth2/v2.0/token", aad_server.uri())); + let ctx = BridgeContext::new( + "req-azure-1", + sample_model(), + sample_pk_with_aad_secret("https://acme-west.openai.azure.com"), + ); + let req = ChatFormat::new("my-azure-gpt4", vec![ChatMessage::user("hi")]); + let resp = bridge.chat(&req, &ctx).await.unwrap(); + assert_eq!(resp.message.content, "hello via aad"); + + let headers = captured_headers + .lock() + .unwrap() + .clone() + .expect("headers captured"); + assert_eq!( + headers.get("authorization").unwrap(), + "Bearer aad.bearer.test-token", + "AAD path must send Authorization: Bearer " + ); + assert!( + !headers.contains_key("api-key"), + "AAD path must NOT send api-key: header (mutually exclusive auth schemes)" + ); + } + + #[tokio::test] + async fn chat_with_aad_secret_caches_minted_token_across_calls() { + // Three back-to-back chats with the same AAD secret must + // share the cache slot — only one token-mint trip. + let aad_server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/oauth2/v2.0/token")) + .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ + "access_token": "aad.cached-token", + "expires_in": 3600, + "token_type": "Bearer" + }))) + .expect(1) // critical: must be called EXACTLY ONCE across 3 chats + .mount(&aad_server) + .await; + + let azure_server = MockServer::start().await; + Mock::given(method("POST")) + .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ + "id": "x", "object": "chat.completion", "created": 1700000000_i64, + "model": "gpt-4o", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2} + }))) + .mount(&azure_server) + .await; + + let bridge = AzureOpenAiBridge::new() + .with_url_override(mock_chat_url(&azure_server.uri(), "gpt4o-prod")) + .with_aad_token_endpoint_override(format!("{}/oauth2/v2.0/token", aad_server.uri())); + let ctx = BridgeContext::new( + "req-azure-1", + sample_model(), + sample_pk_with_aad_secret("https://acme-west.openai.azure.com"), + ); + let req = ChatFormat::new("my-azure-gpt4", vec![ChatMessage::user("hi")]); + for _ in 0..3 { + bridge.chat(&req, &ctx).await.unwrap(); + } + // wiremock's `expect(1)` on the AAD mock asserts on server + // drop that only one mint trip happened across three chats. + } + + #[tokio::test] + async fn chat_with_aad_secret_aad_4xx_surfaces_before_azure_call() { + // AAD 401 bubbles out as Config (operator-actionable) without + // ever calling the Azure OpenAI endpoint. + let aad_server = MockServer::start().await; + Mock::given(method("POST")) + .respond_with(ResponseTemplate::new(401).set_body_string( + r#"{"error":"invalid_client","error_description":"AADSTS7000215"}"#, + )) + .mount(&aad_server) + .await; + + let azure_server = MockServer::start().await; + Mock::given(method("POST")) + .respond_with(ResponseTemplate::new(200)) + .expect(0) // must NOT be called + .mount(&azure_server) + .await; + + let bridge = AzureOpenAiBridge::new() + .with_url_override(mock_chat_url(&azure_server.uri(), "gpt4o-prod")) + .with_aad_token_endpoint_override(format!("{}/oauth2/v2.0/token", aad_server.uri())); + let ctx = BridgeContext::new( + "req-azure-1", + sample_model(), + sample_pk_with_aad_secret("https://acme-west.openai.azure.com"), + ); + let req = ChatFormat::new("my-azure-gpt4", vec![ChatMessage::user("hi")]); + let err = bridge.chat(&req, &ctx).await.unwrap_err(); + match err { + BridgeError::Config(msg) => { + assert!(msg.contains("invalid_client")); + } + other => panic!("expected Config error, got {other:?}"), + } + } } diff --git a/crates/aisix-provider-azure-openai/src/lib.rs b/crates/aisix-provider-azure-openai/src/lib.rs index da742ae0..025ef1f9 100644 --- a/crates/aisix-provider-azure-openai/src/lib.rs +++ b/crates/aisix-provider-azure-openai/src/lib.rs @@ -4,7 +4,8 @@ //! //! ## Status (issue #302 Phase F) //! -//! - [x] D6.1 — `api-key` header auth (NOT `Authorization: Bearer`) +//! - [x] D6.1 — `api-key` header auth (resource-key scheme; +//! `provider_key.secret` is a verbatim string) //! - [x] D6.2 — Azure URL pattern: //! `https://.openai.azure.com/openai/deployments//chat/completions?api-version=` //! - [x] D6.3 — Deployment-keyed dispatch from `Model.model_name` @@ -16,13 +17,19 @@ //! `OpenAiStreamChunk` parsers ignore unknown fields by default //! (no `deny_unknown_fields`), so the extension passes through //! without breaking decoding. +//! - [x] D6.6 — AAD (Entra ID) Bearer auth as a second auth scheme. +//! `provider_key.secret` autodetects between the resource-key path +//! (verbatim string) and the AAD client_credentials path (JSON +//! `{tenant_id, client_id, client_secret}`) by checking the leading +//! character. The AAD path mints + caches tokens in-process via the +//! client_credentials grant (no JWT signing) and sends +//! `Authorization: Bearer ` instead of the `api-key:` header. +//! Backward-compatible — existing api-key deployments keep working +//! unchanged. //! - [ ] D6.4 — Per-PK `api_version` override. Today the bridge pins //! `AzureUpstreamRef::DEFAULT_API_VERSION` (GA). Follow-up will //! accept an explicit version from `provider_key.api_base` query //! string or a dedicated PK field. -//! - [ ] D6.6 — AAD Bearer auth as a second auth scheme. Today the -//! bridge supports api-key only (the common case). AAD support will -//! land alongside the cp-api `auth_scheme` field becoming routable. //! //! # Why Azure-OpenAI is a separate `Adapter::AzureOpenai` family bridge //! @@ -61,6 +68,7 @@ #![forbid(unsafe_code)] #![deny(rust_2018_idioms)] +mod aad_token_mint; mod bridge; mod wire; From 9fac7e5f2e5dbc0775d399d478d599ead5eb11d6 Mon Sep 17 00:00:00 2001 From: Ming Wen Date: Sun, 24 May 2026 20:16:23 +0800 Subject: [PATCH 2/2] test(azure-openai): add chat_stream AAD bearer regression guard (audit LOW on #388) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit audit-aigw-388-azure-aad flagged that chat_stream() calls the same resolve_auth helper as chat() but had no test pinning the AAD → Authorization: Bearer flow on the streaming path. A future refactor that accidentally skipped resolve_auth in chat_stream (e.g. moved auth resolution into the chat() future and forgot to mirror it on the stream side) would slip past every existing test. Mirrors the same gap noted in audit-aigw-387 (Vertex SA OAuth) which was deferred there as non-blocking; applying the equivalent guard here while the cost is one short test function. The test pins: - Authorization: Bearer set on the upstream stream request - api-key: NOT set (mutex with bearer path) - Accept: text/event-stream set (matches existing chat_stream contract regardless of auth scheme) cargo test -p aisix-provider-azure-openai → 54/54 PASS (was 53; +1). --- .../aisix-provider-azure-openai/src/bridge.rs | 67 +++++++++++++++++++ 1 file changed, 67 insertions(+) diff --git a/crates/aisix-provider-azure-openai/src/bridge.rs b/crates/aisix-provider-azure-openai/src/bridge.rs index 5f887ef0..e771a49a 100644 --- a/crates/aisix-provider-azure-openai/src/bridge.rs +++ b/crates/aisix-provider-azure-openai/src/bridge.rs @@ -2021,4 +2021,71 @@ mod tests { other => panic!("expected Config error, got {other:?}"), } } + + /// Audit LOW (audit-aigw-388-azure-aad): chat_stream must also + /// resolve auth via the AAD path. The streaming code calls the + /// same `resolve_auth` helper as chat(), so this regression- + /// guards a future refactor that accidentally skipped it. + /// Mirrors the same gap noted (and addressed in follow-up) by + /// audit-aigw-387 on the Vertex SA OAuth side. + #[tokio::test] + async fn chat_stream_with_aad_secret_sets_authorization_bearer_header() { + let aad_server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/oauth2/v2.0/token")) + .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ + "access_token": "aad.stream-bearer", + "expires_in": 3600, + "token_type": "Bearer" + }))) + .expect(1) + .mount(&aad_server) + .await; + + let azure_server = MockServer::start().await; + let captured_headers: std::sync::Arc>> = + std::sync::Arc::new(std::sync::Mutex::new(None)); + let captured_for_responder = captured_headers.clone(); + Mock::given(method("POST")) + .respond_with(move |req: &wiremock::Request| { + *captured_for_responder.lock().unwrap() = Some(req.headers.clone()); + ResponseTemplate::new(200) + .insert_header("content-type", "text/event-stream") + .set_body_string("data: [DONE]\n\n") + }) + .mount(&azure_server) + .await; + + let bridge = AzureOpenAiBridge::new() + .with_url_override(mock_chat_url(&azure_server.uri(), "gpt4o-prod")) + .with_aad_token_endpoint_override(format!("{}/oauth2/v2.0/token", aad_server.uri())); + let ctx = BridgeContext::new( + "req-azure-1", + sample_model(), + sample_pk_with_aad_secret("https://acme-west.openai.azure.com"), + ); + let req = ChatFormat::new("my-azure-gpt4", vec![ChatMessage::user("hi")]); + let mut stream = bridge.chat_stream(&req, &ctx).await.unwrap(); + while stream.next().await.is_some() {} + + let headers = captured_headers + .lock() + .unwrap() + .clone() + .expect("headers captured"); + assert_eq!( + headers.get("authorization").unwrap(), + "Bearer aad.stream-bearer", + "chat_stream AAD path must send Authorization: Bearer " + ); + assert!( + !headers.contains_key("api-key"), + "chat_stream AAD path must NOT send api-key: header" + ); + assert_eq!( + headers.get("accept").unwrap(), + "text/event-stream", + "chat_stream must request SSE via Accept header" + ); + } }