Merge upstream main into feat/500-api-key-ip-whitelist

apps/aether-gateway/src/control/auth/credentials.rs

@@ -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();

apps/aether-gateway/src/control/auth/gate.rs

@@ -95,6 +95,7 @@ pub(crate) async fn request_model_local_rejection(
         decision,
         auth_context,
         requested_model.as_deref(),
+        headers,
         body,
     )
     .await
@@ -105,6 +106,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 {
@@ -147,7 +149,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);
     };
@@ -174,6 +177,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
@@ -184,7 +188,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)

apps/aether-gateway/src/control/route/admin/operations_families.rs

@@ -8,6 +8,56 @@ pub(super) fn classify_admin_operations_family_route(
     normalized_path_no_trailing: &str,
 ) -> Option<ClassifiedRoute> {
     if method == http::Method::GET
+        && matches!(
+            normalized_path,
+            "/api/admin/referrals" | "/api/admin/referrals/"
+        )
+    {
+        Some(classified(
+            "admin_proxy",
+            "referrals_manage",
+            "list_referrals",
+            "admin:billing",
+            false,
+        ))
+    } else if method == http::Method::GET
+        && matches!(
+            normalized_path,
+            "/api/admin/referral-rewards" | "/api/admin/referral-rewards/"
+        )
+    {
+        Some(classified(
+            "admin_proxy",
+            "referrals_manage",
+            "list_referral_rewards",
+            "admin:billing",
+            false,
+        ))
+    } else if method == http::Method::POST
+        && normalized_path.starts_with("/api/admin/referral-rewards/")
+        && normalized_path.ends_with("/retry")
+        && normalized_path.matches('/').count() == 5
+    {
+        Some(classified(
+            "admin_proxy",
+            "referrals_manage",
+            "retry_referral_reward",
+            "admin:billing",
+            false,
+        ))
+    } else if method == http::Method::POST
+        && normalized_path.starts_with("/api/admin/referral-rewards/")
+        && normalized_path.ends_with("/void")
+        && normalized_path.matches('/').count() == 5
+    {
+        Some(classified(
+            "admin_proxy",
+            "referrals_manage",
+            "void_referral_reward",
+            "admin:billing",
+            false,
+        ))
+    } else if method == http::Method::GET
         && matches!(
             normalized_path,
             "/api/admin/provider-ops/architectures" | "/api/admin/provider-ops/architectures/"

apps/aether-gateway/src/control/route/ai.rs

@@ -55,7 +55,7 @@ pub(super) fn classify_ai_public_route(
     } else if method == http::Method::POST
         && matches!(
             normalized_path,
-            "/v1/images/generations" | "/v1/images/edits" | "/v1/images/variations"
+            "/v1/images/generations" | "/v1/images/edits"
         )
     {
         Some(classified(
@@ -104,6 +104,17 @@ pub(super) fn classify_ai_public_route(
                 "gemini:video",
                 true,
             ))
+        } else if normalized_path.ends_with(":embedContent")
+            || normalized_path.ends_with(":batchEmbedContents")
+        {
+            Some(classified_with_request_auth_channel(
+                "ai_public",
+                "gemini",
+                "embedding",
+                "api_key",
+                "gemini:embedding",
+                true,
+            ))
         } else if is_gemini_cli_request(headers) {
             Some(classified_with_request_auth_channel(
                 "ai_public",

apps/aether-gateway/src/control/route/mod.rs

@@ -229,6 +229,8 @@ pub(super) fn is_gemini_models_route(path: &str) -> bool {
     (path.starts_with("/v1/models/") || path.starts_with("/v1beta/models/"))
         && (path.contains(":generateContent")
             || path.contains(":streamGenerateContent")
+            || path.contains(":embedContent")
+            || path.contains(":batchEmbedContents")
             || path.contains(":predictLongRunning"))
 }
 

apps/aether-gateway/src/control/route/public_support.rs

@@ -279,6 +279,20 @@ pub(super) fn classify_public_support_route(
             "user:announcements",
             false,
         ))
+    } else if method == http::Method::GET
+        && matches!(
+            normalized_path,
+            "/api/announcements/users/me/required-unread"
+                | "/api/announcements/users/me/required-unread/"
+        )
+    {
+        Some(classified(
+            "public_support",
+            "announcement_user",
+            "required_unread",
+            "user:announcements",
+            false,
+        ))
     } else if method == http::Method::POST
         && matches!(
             normalized_path,
@@ -462,6 +476,7 @@ pub(super) fn classify_public_support_route(
                 | "/api/users/me/available-models"
                 | "/api/users/me/endpoint-status"
                 | "/api/users/me/preferences"
+                | "/api/users/me/referral"
                 | "/api/users/me/model-capabilities"
         )
     {
@@ -477,6 +492,7 @@ pub(super) fn classify_public_support_route(
             "/api/users/me/available-models" => "available_models",
             "/api/users/me/endpoint-status" => "endpoint_status",
             "/api/users/me/preferences" => "preferences",
+            "/api/users/me/referral" => "referral",
             "/api/users/me/model-capabilities" => "model_capabilities",
             _ => "detail",
         };
@@ -739,6 +755,7 @@ pub(super) fn classify_public_support_route(
         ))
     } else if method == http::Method::GET
         && (has_single_segment_after_prefix(normalized_path, "/install/")
+            || has_single_segment_after_prefix(normalized_path, "/install-tunnel/")
             || has_single_segment_after_prefix(normalized_path, "/install-proxy/")
             || has_single_segment_after_prefix(normalized_path, "/i/"))
     {

apps/aether-gateway/src/control/tests/ai.rs

@@ -80,6 +80,28 @@ fn classifies_openai_chat_and_responses_separately_from_embedding() {
     assert_ne!(responses.route_kind.as_deref(), Some("embedding"));
 }
 
+#[test]
+fn classifies_openai_image_generation_and_edit_but_not_variation() {
+    let headers = headers(&[("authorization", "Bearer sk-test")]);
+
+    for path in ["/v1/images/generations", "/v1/images/edits"] {
+        let uri: Uri = path.parse().expect("uri should parse");
+        let decision = classify_control_route(&http::Method::POST, &uri, &headers)
+            .expect("image route should classify");
+
+        assert_eq!(decision.route_family.as_deref(), Some("openai"));
+        assert_eq!(decision.route_kind.as_deref(), Some("image"));
+        assert_eq!(
+            decision.auth_endpoint_signature.as_deref(),
+            Some("openai:image")
+        );
+        assert!(decision.is_execution_runtime_candidate());
+    }
+
+    let variation_uri: Uri = "/v1/images/variations".parse().expect("uri should parse");
+    assert!(classify_control_route(&http::Method::POST, &variation_uri, &headers).is_none());
+}
+
 #[test]
 fn classifies_models_list_as_claude_when_headers_match() {
     let headers = headers(&[
@@ -197,6 +219,44 @@ fn classifies_gemini_generate_content_api_key_without_cli_marker() {
     assert!(decision.is_execution_runtime_candidate());
 }
 
+#[test]
+fn classifies_gemini_embed_content_as_embedding_route() {
+    let headers = headers(&[("x-goog-api-key", "gemini-key")]);
+    let uri: Uri = "/v1beta/models/gemini-embedding-2-preview:embedContent"
+        .parse()
+        .expect("uri should parse");
+    let decision =
+        classify_control_route(&http::Method::POST, &uri, &headers).expect("route should classify");
+
+    assert_eq!(decision.route_family.as_deref(), Some("gemini"));
+    assert_eq!(decision.route_kind.as_deref(), Some("embedding"));
+    assert_eq!(decision.request_auth_channel.as_deref(), Some("api_key"));
+    assert_eq!(
+        decision.auth_endpoint_signature.as_deref(),
+        Some("gemini:embedding")
+    );
+    assert!(decision.is_execution_runtime_candidate());
+}
+
+#[test]
+fn classifies_gemini_batch_embed_contents_as_embedding_route() {
+    let headers = headers(&[("x-goog-api-key", "gemini-key")]);
+    let uri: Uri = "/v1beta/models/gemini-embedding-2-preview:batchEmbedContents"
+        .parse()
+        .expect("uri should parse");
+    let decision =
+        classify_control_route(&http::Method::POST, &uri, &headers).expect("route should classify");
+
+    assert_eq!(decision.route_family.as_deref(), Some("gemini"));
+    assert_eq!(decision.route_kind.as_deref(), Some("embedding"));
+    assert_eq!(decision.request_auth_channel.as_deref(), Some("api_key"));
+    assert_eq!(
+        decision.auth_endpoint_signature.as_deref(),
+        Some("gemini:embedding")
+    );
+    assert!(decision.is_execution_runtime_candidate());
+}
+
 #[test]
 fn classifies_gemini_predict_long_running_as_video_route() {
     let headers = headers(&[]);