feat(gateway): normalize compressed request bodies

This commit is contained in:
mayrain
2026-05-19 22:57:45 +08:00
parent 57655bdb25
commit 66a54cc39e
9 changed files with 679 additions and 43 deletions

View File

@@ -7,7 +7,7 @@ use url::form_urlencoded;
use crate::{
ai_serving::extract_gemini_model_from_path,
headers::{header_value_str, is_json_request},
headers::{decoded_request_body_bytes, header_value_str, is_json_request},
};
use super::super::GatewayControlDecision;
@@ -31,7 +31,8 @@ pub(crate) fn extract_requested_model(
if !is_json_request(headers) || body.is_empty() {
return None;
}
serde_json::from_slice::<serde_json::Value>(body)
let body = decoded_request_body_bytes(headers, body.as_ref()).ok()?;
serde_json::from_slice::<serde_json::Value>(body.as_ref())
.ok()
.and_then(|payload| {
payload
@@ -356,15 +357,50 @@ pub(super) fn current_unix_secs() -> u64 {
#[cfg(test)]
mod tests {
use super::{
build_auth_context_cache_key, extract_request_credentials, GatewayCredentialCarrier,
GatewayPrimaryCredential, GatewayTrustedAdminHeaders, GatewayTrustedAuthHeaders,
build_auth_context_cache_key, extract_request_credentials, extract_requested_model,
GatewayCredentialCarrier, GatewayPrimaryCredential, GatewayTrustedAdminHeaders,
GatewayTrustedAuthHeaders,
};
use crate::control::GatewayControlDecision;
use axum::body::Bytes;
use axum::http::{self, Uri};
fn uri(path: &str) -> Uri {
path.parse().expect("uri should parse")
}
#[test]
fn extract_requested_model_reads_zstd_encoded_json_body() {
let decision = GatewayControlDecision::synthetic(
"/v1/responses",
Some("ai_public".to_string()),
Some("openai".to_string()),
Some("responses".to_string()),
Some("openai:responses".to_string()),
);
let mut headers = http::HeaderMap::new();
headers.insert(
http::header::CONTENT_TYPE,
http::HeaderValue::from_static("application/json"),
);
headers.insert(
http::header::CONTENT_ENCODING,
http::HeaderValue::from_static("zstd"),
);
let encoded =
zstd::stream::encode_all(br#"{"model":"gpt-5.4","input":"hello"}"#.as_slice(), 0)
.expect("zstd body should encode");
let requested_model = extract_requested_model(
&decision,
&uri("/v1/responses"),
&headers,
&Bytes::from(encoded),
);
assert_eq!(requested_model.as_deref(), Some("gpt-5.4"));
}
#[test]
fn selects_openai_bearer_as_provider_api_key() {
let mut headers = http::HeaderMap::new();

View File

@@ -94,6 +94,7 @@ pub(crate) async fn request_model_local_rejection(
decision,
auth_context,
requested_model.as_deref(),
headers,
body,
)
.await
@@ -104,6 +105,7 @@ async fn balance_capacity_rejection(
decision: &GatewayControlDecision,
auth_context: &GatewayControlAuthContext,
requested_model: Option<&str>,
headers: &http::HeaderMap,
body: &Bytes,
) -> Result<Option<GatewayLocalAuthRejection>, GatewayError> {
if auth_context.api_key_is_standalone {
@@ -146,7 +148,8 @@ async fn balance_capacity_rejection(
return Ok(None);
};
let Some(estimated_cost_usd) =
estimate_request_cost_upper_bound_usd(state, decision, requested_model, body).await?
estimate_request_cost_upper_bound_usd(state, decision, requested_model, headers, body)
.await?
else {
return Ok(None);
};
@@ -173,6 +176,7 @@ async fn estimate_request_cost_upper_bound_usd(
state: &AppState,
decision: &GatewayControlDecision,
requested_model: &str,
headers: &http::HeaderMap,
body: &Bytes,
) -> Result<Option<f64>, GatewayError> {
let Some(api_format) = decision
@@ -183,7 +187,11 @@ async fn estimate_request_cost_upper_bound_usd(
else {
return Ok(None);
};
let body_json = serde_json::from_slice::<serde_json::Value>(body).ok();
let body = crate::headers::decoded_request_body_bytes(headers, body.as_ref()).ok();
let Some(body) = body else {
return Ok(None);
};
let body_json = serde_json::from_slice::<serde_json::Value>(body.as_ref()).ok();
let Some(input_tokens) = body_json
.as_ref()
.map(estimate_json_tokens)