完善 Antigravity 与 Gemini 跨格式兼容

2026-10-10 03:09:50 +08:00 · 2026-07-08 22:34:29 +08:00
parent b1ed6b24b0
commit f9c8ec41f4
8 changed files with 961 additions and 175 deletions
@@ -0,0 +1,169 @@
+use std::collections::BTreeMap;
+use std::sync::Arc;
+
+use serde_json::Value;
+
+use crate::ai_serving::transport::antigravity::{
+    build_antigravity_safe_v1internal_request, build_antigravity_static_identity_headers,
+    classify_local_antigravity_request_support, AntigravityEnvelopeRequestType,
+    AntigravityRequestAuth, AntigravityRequestAuthUnsupportedReason,
+    AntigravityRequestEnvelopeSupport, AntigravityRequestSideSupport,
+    AntigravityRequestSideUnsupportedReason,
+};
+use crate::ai_serving::transport::{
+    build_standard_provider_request_headers, GatewayProviderTransportSnapshot,
+    StandardProviderRequestHeaders, StandardProviderRequestHeadersInput,
+};
+use crate::AppState;
+
+pub(crate) const ANTIGRAVITY_V1INTERNAL_ENVELOPE_NAME: &str = "antigravity:v1internal";
+
+pub(crate) enum AntigravityV1InternalRequestError {
+    TransportUnsupported,
+    EnvelopeUnsupported,
+    UpstreamUrlUnavailable,
+    HeaderRulesApplyFailed,
+}
+
+pub(crate) struct AntigravityV1InternalRequestInput<'a> {
+    pub(crate) state: &'a AppState,
+    pub(crate) parts: &'a http::request::Parts,
+    pub(crate) transport: &'a Arc<GatewayProviderTransportSnapshot>,
+    pub(crate) trace_id: &'a str,
+    pub(crate) mapped_model: &'a str,
+    pub(crate) provider_api_format: &'a str,
+    pub(crate) auth_header: &'a str,
+    pub(crate) auth_value: &'a str,
+    pub(crate) request_headers: &'a http::HeaderMap,
+    pub(crate) original_request_body: &'a Value,
+    pub(crate) gemini_request_body: &'a Value,
+    pub(crate) upstream_is_stream: bool,
+    pub(crate) same_format: bool,
+}
+
+pub(crate) struct AntigravityV1InternalRequest {
+    pub(crate) transport: Arc<GatewayProviderTransportSnapshot>,
+    pub(crate) body: Value,
+    pub(crate) headers: StandardProviderRequestHeaders,
+    pub(crate) upstream_url: String,
+}
+
+pub(crate) async fn build_antigravity_v1internal_provider_request(
+    input: AntigravityV1InternalRequestInput<'_>,
+) -> Result<AntigravityV1InternalRequest, AntigravityV1InternalRequestError> {
+    let payload = build_antigravity_v1internal_payload(
+        input.state,
+        input.transport,
+        input.trace_id,
+        input.mapped_model,
+        input.gemini_request_body,
+    )
+    .await?;
+
+    let upstream_url = crate::ai_serving::build_provider_transport_request_url_for_request_body(
+        &payload.transport,
+        input.provider_api_format,
+        Some(input.mapped_model),
+        input.upstream_is_stream,
+        input.parts.uri.query(),
+        None,
+        Some(&payload.body),
+    )
+    .ok_or(AntigravityV1InternalRequestError::UpstreamUrlUnavailable)?;
+
+    let extra_headers: BTreeMap<String, String> =
+        build_antigravity_static_identity_headers(&payload.auth);
+    let mut headers =
+        build_standard_provider_request_headers(StandardProviderRequestHeadersInput {
+            transport: &payload.transport,
+            provider_api_format: input.provider_api_format,
+            same_format: input.same_format,
+            headers: input.request_headers,
+            auth_header: input.auth_header,
+            auth_value: input.auth_value,
+            extra_headers: &extra_headers,
+            header_rules: payload.transport.endpoint.header_rules.as_ref(),
+            provider_request_body: &payload.body,
+            original_request_body: input.original_request_body,
+            upstream_is_stream: input.upstream_is_stream,
+        })
+        .ok_or(AntigravityV1InternalRequestError::HeaderRulesApplyFailed)?;
+    headers
+        .headers
+        .insert("accept".to_string(), "text/event-stream".to_string());
+
+    Ok(AntigravityV1InternalRequest {
+        transport: payload.transport,
+        body: payload.body,
+        headers,
+        upstream_url,
+    })
+}
+
+struct AntigravityV1InternalPayload {
+    transport: Arc<GatewayProviderTransportSnapshot>,
+    auth: AntigravityRequestAuth,
+    body: Value,
+}
+
+async fn build_antigravity_v1internal_payload(
+    state: &AppState,
+    transport: &Arc<GatewayProviderTransportSnapshot>,
+    trace_id: &str,
+    mapped_model: &str,
+    gemini_request_body: &Value,
+) -> Result<AntigravityV1InternalPayload, AntigravityV1InternalRequestError> {
+    let mut resolved_transport = Arc::clone(transport);
+    let mut antigravity_support = classify_local_antigravity_request_support(
+        &resolved_transport,
+        gemini_request_body,
+        AntigravityEnvelopeRequestType::Agent,
+    );
+
+    if matches!(
+        antigravity_support,
+        AntigravityRequestSideSupport::Unsupported(
+            AntigravityRequestSideUnsupportedReason::UnsupportedAuth(
+                AntigravityRequestAuthUnsupportedReason::MissingProjectId
+            )
+        )
+    ) {
+        if let Some(hydrated) = state
+            .hydrate_antigravity_project_metadata_for_transport(&resolved_transport)
+            .await
+        {
+            resolved_transport = Arc::new(hydrated);
+            antigravity_support = classify_local_antigravity_request_support(
+                &resolved_transport,
+                gemini_request_body,
+                AntigravityEnvelopeRequestType::Agent,
+            );
+        }
+    }
+
+    let auth = match antigravity_support {
+        AntigravityRequestSideSupport::Supported(spec) => spec.auth,
+        AntigravityRequestSideSupport::Unsupported(_) => {
+            return Err(AntigravityV1InternalRequestError::TransportUnsupported);
+        }
+    };
+
+    let body = match build_antigravity_safe_v1internal_request(
+        &auth,
+        trace_id,
+        mapped_model,
+        gemini_request_body,
+        AntigravityEnvelopeRequestType::Agent,
+    ) {
+        AntigravityRequestEnvelopeSupport::Supported(envelope) => envelope,
+        AntigravityRequestEnvelopeSupport::Unsupported(_) => {
+            return Err(AntigravityV1InternalRequestError::EnvelopeUnsupported);
+        }
+    };
+
+    Ok(AntigravityV1InternalPayload {
+        transport: resolved_transport,
+        auth,
+        body,
+    })
+}