feat(gateway): harden provider request execution

Preserve exact request payloads and model client surface and API operation explicitly. Add Anthropic compatibility profiles, bounded stream commitment, and scoped OAuth retry behavior across provider transports.
2026-10-09 10:57:03 +08:00 · 2026-07-27 09:36:31 +08:00
parent 79b70f7b5c
commit 531cf11025
152 changed files with 13984 additions and 2075 deletions
@@ -81,6 +81,8 @@ pub(crate) async fn resolve_local_same_format_provider_decision_input(

    let mut input = build_local_requested_model_decision_input(resolved_input, requested_model);
    input.request_auth_channel = decision.request_auth_channel.clone();
+    input.client_surface = decision.client_surface;
+    input.gateway_credential_carrier = decision.gateway_credential_carrier;
    input.client_session_affinity = client_session_affinity_from_api_request(
        spec_metadata.api_format,
        &parts.headers,
@@ -128,7 +130,7 @@ pub(crate) async fn materialize_local_same_format_provider_candidate_attempts(
        .base_model()
        .unwrap_or(&input.requested_model);
    let (candidates, preselection_skipped) = planner_state
-        .list_selectable_candidates_with_skip_reasons(
+        .list_selectable_candidates_with_skip_reasons_for_request_operation(
            spec_metadata.api_format,
            routing_model,
            spec_metadata.require_streaming,
@@ -137,6 +139,7 @@ pub(crate) async fn materialize_local_same_format_provider_candidate_attempts(
            input.client_session_affinity.as_ref(),
            current_unix_secs(),
            false,
+            spec.operation.map(|operation| operation.as_str()),
        )
        .await?;
    let outcome = materialize_local_execution_candidates_with_serving(
@@ -232,7 +235,7 @@ pub(crate) async fn build_local_same_format_provider_candidate_attempt_source<'a
        .base_model()
        .unwrap_or(&input.requested_model);
    let (candidates, preselection_skipped) = planner_state
-        .list_selectable_candidates_with_skip_reasons(
+        .list_selectable_candidates_with_skip_reasons_for_request_operation(
            spec_metadata.api_format,
            routing_model,
            spec_metadata.require_streaming,
@@ -241,6 +244,7 @@ pub(crate) async fn build_local_same_format_provider_candidate_attempt_source<'a
            input.client_session_affinity.as_ref(),
            current_unix_secs(),
            false,
+            spec.operation.map(|operation| operation.as_str()),
        )
        .await?;

@@ -1,5 +1,8 @@
 use serde_json::json;

+use aether_ai_serving::{AdaptationMode, AiRequestGzipPolicy, OriginalRequestPayload};
+use aether_contracts::{ExecutionResponseBodyMode, EXECUTION_RESPONSE_BODY_MODE_HEADER};
+
 use crate::ai_serving::ai_local_execution_contract_for_formats;
 use crate::ai_serving::build_request_trace_proxy_value;
 use crate::ai_serving::planner::candidate_materialization::{
@@ -61,6 +64,8 @@ pub(crate) async fn maybe_build_local_same_format_provider_decision_payload_for_
    else {
        return Ok(None);
    };
+    let request_redacted = resolved.request_redacted;
+    let compatibility_edits_empty = resolved.compatibility_edits.is_empty();
    let original_request_body_json = if resolved.request_redacted {
        Some(&resolved.provider_request_body)
    } else {
@@ -82,6 +87,51 @@ pub(crate) async fn maybe_build_local_same_format_provider_decision_payload_for_
        .clone()
        .or_else(|| resolve_transport_profile(&resolved.transport));
    let mut extra_fields = serde_json::Map::new();
+    extra_fields.insert(
+        "provider_type".to_string(),
+        json!(resolved.transport.provider.provider_type.as_str()),
+    );
+    if let Some(operation) = spec.operation {
+        extra_fields.insert("api_operation".to_string(), json!(operation.as_str()));
+    }
+    if let Some(client_surface) = input.client_surface {
+        extra_fields.insert("client_surface".to_string(), json!(client_surface.as_str()));
+    }
+    if let Some(carrier) = input.gateway_credential_carrier {
+        extra_fields.insert(
+            "gateway_credential_carrier".to_string(),
+            json!(carrier.as_str()),
+        );
+    }
+    extra_fields.insert(
+        "upstream_credential_mode".to_string(),
+        json!(resolved.transport.key.auth_type.trim().to_ascii_lowercase()),
+    );
+    let mut adaptation_mode = if resolved.compatibility_edits.is_empty() {
+        AdaptationMode::NativeTransparent
+    } else {
+        AdaptationMode::SameFormatCompat
+    };
+    if crate::ai_serving::normalize_api_format_alias(&resolved.provider_api_format)
+        == "claude:messages"
+    {
+        let compatibility_profile =
+            crate::ai_serving::transport::resolve_anthropic_compatibility_profile(
+                &resolved.transport,
+                &resolved.provider_api_format,
+            );
+        extra_fields.insert(
+            "anthropic_compatibility_profile".to_string(),
+            json!(compatibility_profile.as_str()),
+        );
+        if compatibility_profile.uses_claude_code_compatibility() {
+            adaptation_mode = AdaptationMode::SameFormatCompat;
+        }
+    }
+    extra_fields.insert(
+        "adaptation_mode".to_string(),
+        json!(adaptation_mode.as_str()),
+    );
    if let Some(proxy_value) =
        build_request_trace_proxy_value(Some(&resolved.transport), proxy.as_ref())
    {
@@ -227,9 +277,79 @@ pub(crate) async fn maybe_build_local_same_format_provider_decision_payload_for_
        &mut decision,
        Some(transport.as_ref()),
    )?;
+    enforce_provider_api_operation_invariants(
+        spec.operation,
+        decision.provider_request_body.as_mut(),
+        &mut decision.provider_request_headers,
+    );
+    decision.provider_request_body_base64 = original_request_body_base64(
+        parts,
+        decision.provider_request_body.as_ref(),
+        adaptation_mode,
+        request_redacted,
+        compatibility_edits_empty,
+        decision.content_encoding.as_deref(),
+        decision.request_gzip.as_ref(),
+    );
+    decision
+        .provider_request_headers
+        .retain(|name, _| !name.eq_ignore_ascii_case(EXECUTION_RESPONSE_BODY_MODE_HEADER));
+    if !spec_metadata.require_streaming && decision.provider_request_body_base64.is_some() {
+        decision.provider_request_headers.insert(
+            EXECUTION_RESPONSE_BODY_MODE_HEADER.to_string(),
+            ExecutionResponseBodyMode::PreserveBytes
+                .as_str()
+                .to_string(),
+        );
+    }
    Ok(Some(decision))
 }

+fn enforce_provider_api_operation_invariants(
+    operation: Option<crate::ai_serving::ApiOperation>,
+    provider_request_body: Option<&mut serde_json::Value>,
+    provider_request_headers: &mut std::collections::BTreeMap<String, String>,
+) {
+    if operation != Some(crate::ai_serving::ApiOperation::ClaudeCountTokens) {
+        return;
+    }
+
+    if let Some(provider_request_body) = provider_request_body {
+        crate::ai_serving::transport::enforce_same_format_provider_api_operation_body_policy(
+            provider_request_body,
+            operation,
+        );
+    }
+    for header_name in ["accept", "content-type"] {
+        provider_request_headers.retain(|name, _| !name.eq_ignore_ascii_case(header_name));
+        provider_request_headers.insert(header_name.to_string(), "application/json".to_string());
+    }
+}
+
+fn original_request_body_base64(
+    parts: &http::request::Parts,
+    provider_request_body: Option<&serde_json::Value>,
+    adaptation_mode: AdaptationMode,
+    request_redacted: bool,
+    compatibility_edits_empty: bool,
+    content_encoding: Option<&str>,
+    request_gzip: Option<&AiRequestGzipPolicy>,
+) -> Option<String> {
+    if adaptation_mode != AdaptationMode::NativeTransparent
+        || request_redacted
+        || !compatibility_edits_empty
+        || content_encoding.is_some_and(|value| !value.trim().is_empty())
+        || request_gzip.is_some_and(|policy| policy.enabled != Some(false))
+    {
+        return None;
+    }
+
+    parts
+        .extensions
+        .get::<OriginalRequestPayload>()?
+        .body_bytes_base64_if_unchanged(provider_request_body?)
+}
+
 pub(super) async fn mark_skipped_local_same_format_provider_candidate(
    state: &AppState,
    input: &LocalSameFormatProviderDecisionInput,
@@ -313,3 +433,177 @@ pub(super) async fn mark_skipped_local_same_format_provider_candidate_with_failu
    )
    .await;
 }
+
+#[cfg(test)]
+mod tests {
+    use std::collections::BTreeMap;
+
+    use base64::Engine as _;
+
+    use super::{
+        enforce_provider_api_operation_invariants, original_request_body_base64, AdaptationMode,
+        AiRequestGzipPolicy, OriginalRequestPayload,
+    };
+    use crate::ai_serving::ApiOperation;
+
+    fn request_parts_with_original_payload(
+        body_json: serde_json::Value,
+        body_bytes: &[u8],
+    ) -> http::request::Parts {
+        let (mut parts, ()) = http::Request::new(()).into_parts();
+        parts
+            .extensions
+            .insert(OriginalRequestPayload::from_parsed_json(
+                body_json, body_bytes,
+            ));
+        parts
+    }
+
+    #[test]
+    fn count_tokens_invariants_win_after_provider_routing_mutations() {
+        let mut body = serde_json::json!({
+            "model": "claude-sonnet-4",
+            "messages": [],
+            "stream": true
+        });
+        let mut headers = BTreeMap::from([
+            ("Accept".to_string(), "text/event-stream".to_string()),
+            ("Content-Type".to_string(), "text/plain".to_string()),
+            ("x-provider-route".to_string(), "kept".to_string()),
+        ]);
+
+        enforce_provider_api_operation_invariants(
+            Some(ApiOperation::ClaudeCountTokens),
+            Some(&mut body),
+            &mut headers,
+        );
+
+        assert!(body.get("stream").is_none());
+        assert_eq!(
+            headers.get("accept").map(String::as_str),
+            Some("application/json")
+        );
+        assert_eq!(
+            headers.get("content-type").map(String::as_str),
+            Some("application/json")
+        );
+        assert_eq!(
+            headers.get("x-provider-route").map(String::as_str),
+            Some("kept")
+        );
+        assert_eq!(
+            headers
+                .keys()
+                .filter(|name| name.eq_ignore_ascii_case("accept"))
+                .count(),
+            1
+        );
+        assert_eq!(
+            headers
+                .keys()
+                .filter(|name| name.eq_ignore_ascii_case("content-type"))
+                .count(),
+            1
+        );
+    }
+
+    #[test]
+    fn unchanged_same_format_body_preserves_original_json_bytes() {
+        let raw = br#"{ "unknown": {"enabled":true}, "messages": [], "model": "claude-sonnet-4" }"#;
+        let body_json: serde_json::Value = serde_json::from_slice(raw).expect("body should parse");
+        let parts = request_parts_with_original_payload(body_json.clone(), raw);
+
+        let encoded = original_request_body_base64(
+            &parts,
+            Some(&body_json),
+            AdaptationMode::NativeTransparent,
+            false,
+            true,
+            None,
+            None,
+        )
+        .expect("unchanged request should retain exact bytes");
+
+        assert_eq!(
+            base64::engine::general_purpose::STANDARD
+                .decode(encoded)
+                .expect("body should decode"),
+            raw
+        );
+    }
+
+    #[test]
+    fn request_edits_or_encoding_disable_original_json_bytes() {
+        let raw = br#"{"model":"claude-sonnet-4","messages":[]}"#;
+        let body_json: serde_json::Value = serde_json::from_slice(raw).expect("body should parse");
+        let parts = request_parts_with_original_payload(body_json.clone(), raw);
+        let changed_body = serde_json::json!({
+            "model": "claude-sonnet-4-5",
+            "messages": []
+        });
+
+        assert!(original_request_body_base64(
+            &parts,
+            Some(&changed_body),
+            AdaptationMode::NativeTransparent,
+            false,
+            true,
+            None,
+            None,
+        )
+        .is_none());
+        assert!(original_request_body_base64(
+            &parts,
+            Some(&body_json),
+            AdaptationMode::NativeTransparent,
+            true,
+            true,
+            None,
+            None,
+        )
+        .is_none());
+        assert!(original_request_body_base64(
+            &parts,
+            Some(&body_json),
+            AdaptationMode::NativeTransparent,
+            false,
+            false,
+            None,
+            None,
+        )
+        .is_none());
+        assert!(original_request_body_base64(
+            &parts,
+            Some(&body_json),
+            AdaptationMode::SameFormatCompat,
+            false,
+            true,
+            None,
+            None,
+        )
+        .is_none());
+        assert!(original_request_body_base64(
+            &parts,
+            Some(&body_json),
+            AdaptationMode::NativeTransparent,
+            false,
+            true,
+            Some("gzip"),
+            None,
+        )
+        .is_none());
+        assert!(original_request_body_base64(
+            &parts,
+            Some(&body_json),
+            AdaptationMode::NativeTransparent,
+            false,
+            true,
+            None,
+            Some(&AiRequestGzipPolicy {
+                enabled: Some(true),
+                min_bytes: Some(1),
+            }),
+        )
+        .is_none());
+    }
+}
@@ -75,9 +75,10 @@ pub(crate) fn resolve_same_format_provider_transport_unsupported_reason_for_trac
            decision_kind: "trace_candidate_metadata",
            report_kind: Some("trace_candidate_metadata"),
        },
+        None,
    );
    if !behavior.is_antigravity
-        && !behavior.is_claude_code
+        && !behavior.is_claude_code_transport
        && !behavior.is_gemini_cli
        && !behavior.is_vertex
        && !behavior.is_kiro
@@ -127,6 +128,23 @@ pub(crate) async fn resolve_local_same_format_provider_candidate_payload_parts(
    spec: LocalSameFormatProviderSpec,
 ) -> Result<Option<LocalSameFormatProviderCandidatePayloadParts>, GatewayError> {
    let candidate = &attempt.eligible.candidate;
+    if let Some(skip_reason) = same_format_provider_operation_skip_reason(
+        &attempt.eligible.transport,
+        attempt.eligible.provider_api_format.as_str(),
+        spec.operation,
+    ) {
+        mark_skipped_local_same_format_provider_candidate(
+            state,
+            input,
+            trace_id,
+            candidate,
+            attempt.candidate_index,
+            &attempt.candidate_id,
+            skip_reason,
+        )
+        .await;
+        return Ok(None);
+    }
    let Some(prepared) = prepare_local_same_format_provider_candidate(
        state,
        trace_id,
@@ -364,7 +382,7 @@ pub(crate) async fn resolve_local_same_format_provider_candidate_payload_parts(
    } else {
        None
    };
-    let provider_request_body = if let Some(antigravity_auth) = antigravity_auth.as_ref() {
+    let mut provider_request_body = if let Some(antigravity_auth) = antigravity_auth.as_ref() {
        match build_antigravity_safe_v1internal_request(
            antigravity_auth,
            trace_id,
@@ -424,6 +442,16 @@ pub(crate) async fn resolve_local_same_format_provider_candidate_payload_parts(
    } else {
        base_provider_request_body
    };
+    if crate::ai_serving::transport::enforce_same_format_provider_api_operation_body_policy(
+        &mut provider_request_body,
+        spec.operation,
+    ) {
+        compatibility_edits.push(SameFormatProviderCompatibilityEdit {
+            field: "stream".to_string(),
+            action: SameFormatProviderCompatibilityEditAction::RuntimeRewrite,
+            detail: "removed stream field for non-streaming API operation".to_string(),
+        });
+    }

    let is_grok = prepared
        .transport
@@ -490,10 +518,10 @@ pub(crate) async fn resolve_local_same_format_provider_candidate_payload_parts(
            original_request_body: body_json,
            header_rules: transport.endpoint.header_rules.as_ref(),
            behavior: prepared.behavior,
+            api_operation: spec.operation,
            auth_header: prepared.auth_header.as_deref(),
            auth_value: prepared.auth_value.as_deref(),
            extra_headers: &extra_headers,
-            key_fingerprint: transport.key.fingerprint.as_ref(),
            kiro_auth_config: prepared.kiro_auth.as_ref().map(|auth| &auth.auth_config),
            kiro_machine_id: prepared
                .kiro_auth
@@ -564,3 +592,98 @@ pub(crate) async fn resolve_local_same_format_provider_candidate_payload_parts(
        request_redacted: redaction.redacted,
    }))
 }
+
+fn same_format_provider_operation_skip_reason(
+    transport: &GatewayProviderTransportSnapshot,
+    provider_api_format: &str,
+    operation: Option<crate::ai_serving::ApiOperation>,
+) -> Option<&'static str> {
+    (!crate::ai_serving::transport::transport_supports_api_operation(
+        transport,
+        provider_api_format,
+        operation,
+    ))
+    .then_some("transport_operation_unsupported")
+}
+
+#[cfg(test)]
+mod tests {
+    use super::same_format_provider_operation_skip_reason;
+    use crate::ai_serving::transport::snapshot::{
+        GatewayProviderTransportEndpoint, GatewayProviderTransportKey,
+        GatewayProviderTransportProvider,
+    };
+    use crate::ai_serving::{ApiOperation, GatewayProviderTransportSnapshot};
+
+    fn private_adapter_transport(provider_type: &str) -> GatewayProviderTransportSnapshot {
+        GatewayProviderTransportSnapshot {
+            provider: GatewayProviderTransportProvider {
+                id: "provider-1".to_string(),
+                name: provider_type.to_string(),
+                provider_type: provider_type.to_string(),
+                website: None,
+                is_active: true,
+                keep_priority_on_conversion: false,
+                enable_format_conversion: true,
+                concurrent_limit: None,
+                max_retries: None,
+                proxy: None,
+                request_timeout_secs: None,
+                stream_first_byte_timeout_secs: None,
+                config: None,
+            },
+            endpoint: GatewayProviderTransportEndpoint {
+                id: "endpoint-1".to_string(),
+                provider_id: "provider-1".to_string(),
+                api_format: "claude:messages".to_string(),
+                api_family: Some("claude".to_string()),
+                endpoint_kind: Some("chat".to_string()),
+                is_active: true,
+                base_url: "https://private.example".to_string(),
+                header_rules: None,
+                body_rules: None,
+                max_retries: None,
+                custom_path: None,
+                config: None,
+                format_acceptance_config: None,
+                proxy: None,
+            },
+            key: GatewayProviderTransportKey {
+                id: "key-1".to_string(),
+                provider_id: "provider-1".to_string(),
+                name: "key".to_string(),
+                auth_type: "oauth".to_string(),
+                is_active: true,
+                api_formats: None,
+                auth_type_by_format: None,
+                allow_auth_channel_mismatch_formats: None,
+                allowed_models: None,
+                capabilities: None,
+                rate_multipliers: None,
+                global_priority_by_format: None,
+                expires_at_unix_secs: None,
+                proxy: None,
+                fingerprint: None,
+                upstream_metadata: None,
+                decrypted_api_key: String::new(),
+                decrypted_auth_config: None,
+            },
+        }
+    }
+
+    #[test]
+    fn private_adapter_count_tokens_is_rejected_by_pre_auth_operation_gate() {
+        for provider_type in ["kiro", "grok"] {
+            let transport = private_adapter_transport(provider_type);
+            assert_eq!(
+                same_format_provider_operation_skip_reason(
+                    &transport,
+                    "claude:messages",
+                    Some(ApiOperation::ClaudeCountTokens),
+                ),
+                Some("transport_operation_unsupported"),
+                "provider_type={provider_type}"
+            );
+        }
+    }
+}
@@ -1,6 +1,6 @@
 use crate::ai_serving::planner::spec_metadata::LocalExecutionSurfaceSpecMetadata;
 use crate::ai_serving::transport::{
-    classify_same_format_provider_request_behavior as classify_same_format_provider_request_behavior_impl,
+    classify_same_format_provider_request_behavior_for_operation as classify_same_format_provider_request_behavior_impl,
    resolve_same_format_provider_direct_auth as resolve_same_format_provider_direct_auth_impl,
    same_format_provider_transport_supported as same_format_provider_transport_supported_impl,
    same_format_provider_transport_unsupported_reason as same_format_provider_transport_unsupported_reason_impl,
@@ -15,6 +15,7 @@ pub(super) fn classify_same_format_provider_request_behavior(
    transport: &GatewayProviderTransportSnapshot,
    provider_api_format: &str,
    spec_metadata: LocalExecutionSurfaceSpecMetadata,
+    api_operation: Option<crate::ai_serving::ApiOperation>,
 ) -> SameFormatProviderRequestBehavior {
    classify_same_format_provider_request_behavior_impl(
        transport,
@@ -25,6 +26,7 @@ pub(super) fn classify_same_format_provider_request_behavior(
                .report_kind
                .expect("same-format provider specs should declare report kind"),
        },
+        api_operation,
    )
 }

@@ -59,6 +59,7 @@ pub(super) async fn prepare_local_same_format_provider_candidate(
        &transport,
        provider_api_format,
        spec_metadata,
+        spec.operation,
    );

    if !same_format_provider_transport_supported(
@@ -214,6 +214,16 @@ impl LocalExecutionAttemptSource<AiSyncAttempt> for LocalSameFormatProviderSyncA
        Ok(drained)
    }

+    async fn skip_credential(&mut self, key_id: &str) -> Result<(), GatewayError> {
+        self.candidates.skip_credential(key_id);
+        Ok(())
+    }
+
+    async fn skip_endpoint(&mut self, endpoint_id: &str) -> Result<(), GatewayError> {
+        self.candidates.skip_endpoint(endpoint_id);
+        Ok(())
+    }
+
    async fn skip_provider(&mut self, provider_id: &str) -> Result<(), GatewayError> {
        self.candidates.skip_provider(provider_id);
        Ok(())
@@ -249,6 +259,16 @@ impl LocalExecutionAttemptSource<AiStreamAttempt>
        Ok(drained)
    }

+    async fn skip_credential(&mut self, key_id: &str) -> Result<(), GatewayError> {
+        self.candidates.skip_credential(key_id);
+        Ok(())
+    }
+
+    async fn skip_endpoint(&mut self, endpoint_id: &str) -> Result<(), GatewayError> {
+        self.candidates.skip_endpoint(endpoint_id);
+        Ok(())
+    }
+
    async fn skip_provider(&mut self, provider_id: &str) -> Result<(), GatewayError> {
        self.candidates.skip_provider(provider_id);
        Ok(())
@@ -24,6 +24,7 @@ pub(crate) fn build_same_format_upstream_url(
            upstream_is_stream,
            request_query: parts.uri.query(),
            kiro_api_region: kiro_auth.map(|auth| auth.auth_config.effective_api_region()),
+            api_operation: spec.operation,
            provider_request_body,
        },
    )