fix(proxy): apply per-PK request overrides on the raw-HTTP handlers (responses/rerank/audio)

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<aisix_core::ProviderKey> {
+        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<ApiKey> {
         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);
+    }
 }

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<aisix_core::ProviderKey> {
+        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<ApiKey> {
         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);
+    }
 }