diff --git a/crates/aisix-proxy/src/audio.rs b/crates/aisix-proxy/src/audio.rs index ff3d2a4e..9ca567ff 100644 --- a/crates/aisix-proxy/src/audio.rs +++ b/crates/aisix-proxy/src/audio.rs @@ -376,11 +376,35 @@ async fn multipart_dispatch( form = form.part(name, part); } + // Build headers explicitly so the PK's `request.default_headers` can inject + // operator headers (AISIX-Cloud#867 follow-up). The body is a multipart + // form, so the JSON body-field overrides don't apply here — only headers do. + // Content-Type is left to `.multipart()` (it sets the boundary). Reserved + // auth headers are protected by `apply_default_headers`. + let mut headers = axum::http::HeaderMap::new(); + let auth_hv = header::HeaderValue::from_str(&format!("Bearer {api_key}")).map_err(|e| { + ProxyError::Bridge(aisix_gateway::BridgeError::Config(format!( + "api key contains invalid header chars: {e}" + ))) + })?; + headers.insert(header::AUTHORIZATION, auth_hv); + let rid_hv = header::HeaderValue::from_str(request_id).map_err(|e| { + ProxyError::Bridge(aisix_gateway::BridgeError::Config(format!( + "request_id contains invalid header chars: {e}" + ))) + })?; + headers.insert( + header::HeaderName::from_static("x-aisix-request-id"), + rid_hv, + ); + if let Some(r) = pk_entry.value.request.as_ref() { + aisix_provider_openai::overrides::apply_default_headers(&mut headers, &r.default_headers); + } + let client = crate::http_client::client(); let resp = client .post(&url) - .header(header::AUTHORIZATION, format!("Bearer {api_key}")) - .header("x-aisix-request-id", request_id) + .headers(headers) .multipart(form) .send() .await @@ -542,12 +566,48 @@ async fn speech_dispatch( *m = Value::String(upstream_model); } + // Apply the PK's `request.*` overrides (body + headers) like the OpenAI + // bridge's chat() path — /v1/audio/speech is a JSON passthrough that builds + // the request directly (AISIX-Cloud#867 follow-up). No-op when none set. + if let Some(r) = pk_entry.value.request.as_ref() { + aisix_provider_openai::overrides::apply_param_renames(&mut body, &r.param_renames); + if let Some(constraints) = &r.param_constraints { + aisix_provider_openai::overrides::apply_param_constraints(&mut body, constraints); + } + aisix_provider_openai::overrides::apply_default_body_fields( + &mut body, + &r.default_body_fields, + ); + } + + let mut headers = axum::http::HeaderMap::new(); + let auth_hv = header::HeaderValue::from_str(&format!("Bearer {api_key}")).map_err(|e| { + ProxyError::Bridge(aisix_gateway::BridgeError::Config(format!( + "api key contains invalid header chars: {e}" + ))) + })?; + headers.insert(header::AUTHORIZATION, auth_hv); + headers.insert( + header::CONTENT_TYPE, + header::HeaderValue::from_static("application/json"), + ); + let rid_hv = header::HeaderValue::from_str(request_id).map_err(|e| { + ProxyError::Bridge(aisix_gateway::BridgeError::Config(format!( + "request_id contains invalid header chars: {e}" + ))) + })?; + headers.insert( + header::HeaderName::from_static("x-aisix-request-id"), + rid_hv, + ); + if let Some(r) = pk_entry.value.request.as_ref() { + aisix_provider_openai::overrides::apply_default_headers(&mut headers, &r.default_headers); + } + let client = crate::http_client::client(); let resp = client .post(crate::dispatch::build_v1_url(&base, "/audio/speech")) - .header(header::AUTHORIZATION, format!("Bearer {api_key}")) - .header(header::CONTENT_TYPE, "application/json") - .header("x-aisix-request-id", request_id) + .headers(headers) .json(&body) .send() .await @@ -747,7 +807,7 @@ mod tests { use axum::body::to_bytes; use axum::http::{Request, StatusCode}; use std::sync::Arc; - use wiremock::matchers::{method, path}; + use wiremock::matchers::{body_partial_json, header, method, path}; use wiremock::{Mock, MockServer, ResponseTemplate}; fn cfg() -> ProxyConfig { @@ -818,6 +878,24 @@ mod tests { snap } + /// A PK carrying `request.*` operator overrides (AISIX-Cloud#867): + /// a default body field + a default header that the audio handlers + /// must apply to the upstream request. + fn provider_key_entry_overrides(api_base: &str) -> ResourceEntry { + let json = format!( + r#"{{"display_name":"openai-up","secret":"sk-up","api_base":"{api_base}","provider":"openai","adapter":"openai","request":{{"default_body_fields":{{"safe_flag":true}},"default_headers":{{"x-custom":"trace-on"}}}}}}"# + ); + let pk: aisix_core::ProviderKey = serde_json::from_str(&json).unwrap(); + ResourceEntry::new(PK_ID, pk, 1) + } + + fn new_snap_overrides(api_base: &str) -> AisixSnapshot { + let snap = AisixSnapshot::new(); + snap.provider_keys + .insert(provider_key_entry_overrides(api_base)); + snap + } + fn apikey_entry(allowed: &[&str]) -> ResourceEntry { let json = format!( r#"{{"key_hash": "8b6712790a2089c67aa97a2d80022df18cc65c7814350e33baebe79aab508891", "allowed_models": {}}}"#, @@ -1250,4 +1328,72 @@ mod tests { }); assert_eq!(super::extract_token_usage(&v), None); } + + /// AISIX-Cloud#867: `/v1/audio/speech` (JSON body) must apply the PK's + /// `request.*` overrides to BOTH the request body + /// (`default_body_fields`) and the request headers (`default_headers`). + /// The Mock matches only when the upstream request carries the injected + /// body field AND header, so a 200 proves both were applied. + #[tokio::test] + async fn speech_applies_pk_request_overrides_issue_867() { + let upstream = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/v1/audio/speech")) + .and(body_partial_json(serde_json::json!({"safe_flag": true}))) + .and(header("x-custom", "trace-on")) + .respond_with(ResponseTemplate::new(200).set_body_bytes(b"AUDIO".to_vec())) + .mount(&upstream) + .await; + + let snap = new_snap_overrides(&upstream.uri()); + snap.models.insert(tts_model("my-tts")); + snap.apikeys.insert(apikey_entry(&["*"])); + + let app = build_app(snap); + let req = Request::builder() + .method("POST") + .uri("/v1/audio/speech") + .header("authorization", "Bearer sk-caller") + .header("content-type", "application/json") + .body(axum::body::Body::from( + r#"{"model":"my-tts","input":"hi","voice":"alloy"}"#, + )) + .unwrap(); + let resp = tower::ServiceExt::oneshot(app, req).await.unwrap(); + assert_eq!(resp.status(), StatusCode::OK); + } + + /// AISIX-Cloud#867: `/v1/audio/transcriptions` (multipart body) must + /// apply the PK's `request.default_headers` to the upstream request. + /// Body `request.*` overrides do NOT apply (the body is a multipart + /// form, not JSON). The Mock matches only on the injected header, so a + /// 200 proves the operator header was applied. + #[tokio::test] + async fn transcriptions_applies_default_headers_issue_867() { + let upstream = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/v1/audio/transcriptions")) + .and(header("x-custom", "trace-on")) + .respond_with( + ResponseTemplate::new(200).set_body_json(serde_json::json!({"text": "hi"})), + ) + .mount(&upstream) + .await; + + let snap = new_snap_overrides(&upstream.uri()); + snap.models.insert(whisper_model("my-transcribe")); + snap.apikeys.insert(apikey_entry(&["*"])); + + let app = build_app(snap); + let (ct, body) = transcription_multipart("my-transcribe"); + let req = Request::builder() + .method("POST") + .uri("/v1/audio/transcriptions") + .header("authorization", "Bearer sk-caller") + .header("content-type", ct) + .body(body) + .unwrap(); + let resp = tower::ServiceExt::oneshot(app, req).await.unwrap(); + assert_eq!(resp.status(), StatusCode::OK); + } } diff --git a/crates/aisix-proxy/src/rerank.rs b/crates/aisix-proxy/src/rerank.rs index 75900a88..0f5e0a53 100644 --- a/crates/aisix-proxy/src/rerank.rs +++ b/crates/aisix-proxy/src/rerank.rs @@ -275,6 +275,18 @@ async fn dispatch( *m = Value::String(upstream_model.clone()); } + // Apply the PK's `request.*` body overrides, matching the OpenAI bridge's + // chat() path and /v1/messages passthrough (AISIX-Cloud#867 follow-up). The + // /v1/rerank path builds the request directly, so without this the override + // pipeline silently no-ops here. No-op when the PK carries none. + if let Some(r) = pk_entry.value.request.as_ref() { + aisix_provider_openai::overrides::apply_param_renames(body, &r.param_renames); + if let Some(constraints) = &r.param_constraints { + aisix_provider_openai::overrides::apply_param_constraints(body, constraints); + } + aisix_provider_openai::overrides::apply_default_body_fields(body, &r.default_body_fields); + } + // Build upstream URL. build_v1_url tolerates either base form — // `https://api.cohere.com` (bare host) and `https://api.openai.com/v1` // (OpenAI-SDK convention, with /v1) both end up at `…/v1/rerank` @@ -295,13 +307,34 @@ async fn dispatch( }; let url = crate::dispatch::build_v1_url(&base, "/rerank"); + // Build headers explicitly so the PK's `request.default_headers` can inject + // operator headers (reserved auth headers are protected by the apply step). + let mut headers = axum::http::HeaderMap::new(); + let auth_hv = HeaderValue::from_str(&format!("Bearer {api_key}")).map_err(|e| { + ProxyError::Bridge(aisix_gateway::BridgeError::Config(format!( + "api key contains invalid header chars: {e}" + ))) + })?; + headers.insert(axum::http::header::AUTHORIZATION, auth_hv); + headers.insert( + axum::http::header::CONTENT_TYPE, + HeaderValue::from_static("application/json"), + ); + let rid_hv = HeaderValue::from_str(request_id).map_err(|e| { + ProxyError::Bridge(aisix_gateway::BridgeError::Config(format!( + "request_id contains invalid header chars: {e}" + ))) + })?; + headers.insert( + axum::http::header::HeaderName::from_static("x-aisix-request-id"), + rid_hv, + ); + if let Some(r) = pk_entry.value.request.as_ref() { + aisix_provider_openai::overrides::apply_default_headers(&mut headers, &r.default_headers); + } + let client = crate::http_client::client(); - let mut req = client - .post(&url) - .header("authorization", format!("Bearer {api_key}")) - .header("content-type", "application/json") - .header("x-aisix-request-id", request_id) - .json(body); + let mut req = client.post(&url).headers(headers).json(body); // #554: rerank is non-streaming; apply the E2E request timeout. if let Some(d) = model.request_timeout() { req = req.timeout(d); @@ -632,6 +665,26 @@ mod tests { snap } + /// AISIX-Cloud#867: an OpenAI PK that carries `request.*` overrides + /// (`default_body_fields` + `default_headers`). Clones the plain openai PK + /// JSON and appends a `request` block; reuses `PK_ID` so the rerank model + /// fixtures still reference it. Used to prove the resolved PK's request + /// overrides reach the rerank upstream body + headers. + fn provider_key_entry_overrides(api_base: &str) -> ResourceEntry { + let json = format!( + r#"{{"display_name":"openai-up","secret":"sk-test","api_base":"{api_base}","provider":"openai","adapter":"openai","request":{{"default_body_fields":{{"safe_flag":true}},"default_headers":{{"x-custom":"trace-on"}}}}}}"# + ); + let pk: aisix_core::ProviderKey = serde_json::from_str(&json).unwrap(); + ResourceEntry::new(PK_ID, pk, 1) + } + + fn new_snap_overrides(api_base: &str) -> AisixSnapshot { + let snap = AisixSnapshot::new(); + snap.provider_keys + .insert(provider_key_entry_overrides(api_base)); + snap + } + fn apikey_entry(allowed: &[&str]) -> ResourceEntry { let json = format!( r#"{{"key_hash":"8b6712790a2089c67aa97a2d80022df18cc65c7814350e33baebe79aab508891","allowed_models":{}}}"#, @@ -1424,4 +1477,49 @@ mod tests { "pk_label must mirror telemetry_tags.pk_label", ); } + + /// AISIX-Cloud#867: the resolved ProviderKey's `request.*` overrides + /// (`default_body_fields` + `default_headers`) must be applied to the + /// outbound /v1/rerank request — exactly like the other proxy passthrough + /// endpoints. The mock matcher ONLY accepts the request when BOTH the + /// injected body field (`safe_flag:true`) and the injected header + /// (`x-custom: trace-on`) are present, so a 200 proves the overrides were + /// applied. Pre-fix the rerank handler dropped them → mock unmatched → + /// non-200. + #[tokio::test] + async fn applies_pk_request_overrides_issue_867() { + use wiremock::matchers::{body_partial_json, header}; + + let upstream = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/v1/rerank")) + .and(body_partial_json(serde_json::json!({"safe_flag": true}))) + .and(header("x-custom", "trace-on")) + .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ + "id": "rerank-1", + "results": [{"index": 0, "relevance_score": 0.9}], + "model": "rerank-multilingual-v3.0", + "usage": {"prompt_tokens": 31, "total_tokens": 31} + }))) + .mount(&upstream) + .await; + + let snap = new_snap_overrides(&upstream.uri()); + snap.models.insert(openai_model("rerank-openai")); + snap.apikeys.insert(apikey_entry(&["*"])); + let app = build_app(snap); + + let resp = app + .oneshot(make_req(serde_json::json!({ + "model": "rerank-openai", + "query": "what is the capital of France?", + "documents": ["Paris", "London", "Berlin"] + }))) + .await + .unwrap(); + + // The mock only matches when both the injected body field and header + // are present — a 200 proves the PK request overrides were applied. + assert_eq!(resp.status(), StatusCode::OK); + } } diff --git a/crates/aisix-proxy/src/responses.rs b/crates/aisix-proxy/src/responses.rs index 24455baf..955187f4 100644 --- a/crates/aisix-proxy/src/responses.rs +++ b/crates/aisix-proxy/src/responses.rs @@ -610,6 +610,23 @@ async fn responses_to_target( *m = Value::String(upstream_model.clone()); } + // Apply the PK's `request.*` overrides to the outbound body, matching the + // OpenAI bridge's chat() path and the /v1/messages passthrough. The + // verbatim /v1/responses path builds the request directly (bypassing the + // Hub), so without this the override pipeline silently no-ops for Codex + // traffic (AISIX-Cloud#867 follow-up). Apply order: renames → constraints + // → defaults; each is a no-op when its configured map is empty. + if let Some(r) = pk_entry.value.request.as_ref() { + aisix_provider_openai::overrides::apply_param_renames(&mut body, &r.param_renames); + if let Some(constraints) = &r.param_constraints { + aisix_provider_openai::overrides::apply_param_constraints(&mut body, constraints); + } + aisix_provider_openai::overrides::apply_default_body_fields( + &mut body, + &r.default_body_fields, + ); + } + let base = crate::dispatch::resolve_base_url(&pk_entry.value)?; // build_v1_url tolerates both `https://api.openai.com` (provider // default) and `https://api.openai.com/v1` (the OpenAI-SDK form @@ -623,13 +640,33 @@ async fn responses_to_target( .and_then(|v| v.as_bool()) .unwrap_or(false); + // Build headers explicitly so the PK's `request.default_headers` can inject + // operator headers. Bridge-owned headers go in FIRST; `apply_default_headers` + // skips already-present keys + the reserved auth blacklist, so an operator + // header can never clobber auth. + let mut headers = axum::http::HeaderMap::new(); + let auth_hv = HeaderValue::from_str(&format!("Bearer {api_key}")).map_err(|e| { + ProxyError::Bridge(aisix_gateway::BridgeError::Config(format!( + "api key contains invalid header chars: {e}" + ))) + })?; + headers.insert(axum::http::header::AUTHORIZATION, auth_hv); + headers.insert( + axum::http::header::CONTENT_TYPE, + HeaderValue::from_static("application/json"), + ); + let rid_hv = HeaderValue::from_str(request_id).map_err(|e| { + ProxyError::Bridge(aisix_gateway::BridgeError::Config(format!( + "request_id contains invalid header chars: {e}" + ))) + })?; + headers.insert(HeaderName::from_static("x-aisix-request-id"), rid_hv); + if let Some(r) = pk_entry.value.request.as_ref() { + aisix_provider_openai::overrides::apply_default_headers(&mut headers, &r.default_headers); + } + let client = crate::http_client::client(); - let mut req = client - .post(&url) - .header("authorization", format!("Bearer {api_key}")) - .header("content-type", "application/json") - .header("x-aisix-request-id", request_id) - .json(&body); + let mut req = client.post(&url).headers(headers).json(&body); // #554: non-streaming gets the E2E request timeout via reqwest's // request-level timeout. Streaming must NOT use it (it would cap the // whole stream); the streaming branch below enforces the per-chunk @@ -1799,7 +1836,7 @@ mod tests { use axum::http::{Request, StatusCode}; use std::sync::Arc; use tower::ServiceExt; - use wiremock::matchers::{header, method, path}; + use wiremock::matchers::{body_partial_json, header, method, path}; use wiremock::{Mock, MockServer, ResponseTemplate}; fn cfg() -> ProxyConfig { @@ -1838,6 +1875,18 @@ mod tests { ResourceEntry::new(OPENAI_PK_ID, pk, 1) } + /// An OpenAI PK carrying per-PK `request.*` overrides (AISIX-Cloud#867): + /// a `default_body_fields` injection and a `default_headers` injection, + /// so the verbatim Responses path can be asserted to apply both to the + /// outbound upstream call. + fn openai_pk_with_overrides(api_base: &str) -> ResourceEntry { + let json = format!( + r#"{{"display_name":"openai-up","secret":"sk-test","api_base":"{api_base}","provider":"openai","adapter":"openai","request":{{"default_body_fields":{{"safe_flag":true}},"default_headers":{{"x-custom":"trace-on"}}}}}}"# + ); + let pk: aisix_core::ProviderKey = serde_json::from_str(&json).unwrap(); + ResourceEntry::new(OPENAI_PK_ID, pk, 1) + } + /// An OpenAI PK carrying per-PK telemetry attribution tags /// (AISIX-Cloud#867) so emitted UsageEvents can be asserted to surface the /// upstream vendor + PK label the dashboard's Logs detail shows. @@ -3103,6 +3152,48 @@ mod tests { assert!(!event.occurred_at.is_empty()); } + /// AISIX-Cloud#867: the verbatim-OpenAI /v1/responses path must apply the + /// resolved ProviderKey's `request.*` overrides to the outbound call — + /// both the `default_body_fields` injection (body) and the + /// `default_headers` injection (header) must reach the upstream. The mock + /// only matches (200) when BOTH the injected body field AND header are + /// present, so a 200 proves the overrides were applied. Before the fix the + /// outbound body/headers carried neither → mock wouldn't match → wiremock + /// 404 → non-200. + #[tokio::test] + async fn responses_verbatim_applies_pk_request_overrides_issue_867() { + let upstream = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/v1/responses")) + .and(body_partial_json(serde_json::json!({"safe_flag": true}))) + .and(header("x-custom", "trace-on")) + .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ + "id": "resp-1", + "object": "response", + "output": [], + "usage": {"input_tokens": 3, "output_tokens": 1} + }))) + .mount(&upstream) + .await; + + let snap = AisixSnapshot::new(); + snap.provider_keys + .insert(openai_pk_with_overrides(&upstream.uri())); + snap.models.insert(openai_model("gpt-resp")); + snap.apikeys.insert(apikey_entry(&["*"])); + let app = build_app(snap); + + let resp = app + .oneshot(make_req(serde_json::json!({ + "model": "gpt-resp", + "input": "hi" + }))) + .await + .unwrap(); + + assert_eq!(resp.status(), StatusCode::OK); + } + /// #543: an OUTPUT-blocked /v1/responses still records the billed /// upstream tokens (the provider already charged), marked /// `guardrail_blocked`, with status 422 — NOT a zero-token event. Zeroing