feat(proxy): record tunnel stability metrics

2026-10-09 18:59:50 +08:00 · 2026-05-08 22:03:17 +08:00
parent 9a84a6ff6c
commit a703acd1fe
81 changed files with 7228 additions and 1109 deletions
@@ -11,6 +11,7 @@ use crate::ai_serving::planner::candidate_metadata::{
    LocalExecutionCandidateMetadataParts,
 };
 use crate::ai_serving::planner::candidate_source::{
+    preselect_local_execution_candidates_for_api_formats_with_serving,
    preselect_local_execution_candidates_with_serving, LocalCandidatePreselectionKeyMode,
 };
 use crate::ai_serving::planner::common::extract_requested_model_from_request;
@@ -23,7 +24,8 @@ use crate::ai_serving::planner::materialization_policy::{
 use crate::ai_serving::planner::spec_metadata::local_standard_spec_metadata;
 use crate::ai_serving::{
    ai_local_execution_contract_for_formats, extract_pool_sticky_session_token,
-    resolve_local_decision_execution_runtime_auth_context, GatewayControlDecision, PlannerAppState,
+    gemini_request_is_image_generation, resolve_local_decision_execution_runtime_auth_context,
+    GatewayControlDecision, PlannerAppState,
 };
 use crate::client_session_affinity::client_session_affinity_from_parts;
 use crate::{AppState, GatewayError};
@@ -86,6 +88,8 @@ pub(super) async fn materialize_local_standard_candidate_attempts(
    spec: LocalStandardSpec,
 ) -> Result<(Vec<LocalStandardCandidateAttempt>, usize), GatewayError> {
    let spec_metadata = local_standard_spec_metadata(spec);
+    let is_gemini_image_bridge = spec_metadata.api_format == "gemini:generate_content"
+        && gemini_request_is_image_generation(body_json);
    let planner_state = PlannerAppState::new(state);
    let sticky_session_token = extract_pool_sticky_session_token(body_json);
    let persistence_policy = build_local_candidate_persistence_policy(
@@ -105,6 +109,17 @@ pub(super) async fn materialize_local_standard_candidate_attempts(
        LocalCandidatePreselectionKeyMode::ProviderEndpointKeyModelAndApiFormat,
    )
    .await?;
+    let (candidates, skipped_candidates) = maybe_append_gemini_image_openai_image_preselection(
+        state,
+        planner_state,
+        trace_id,
+        input,
+        body_json,
+        spec,
+        preselection.candidates,
+        preselection.skipped_candidates,
+    )
+    .await?;
    let outcome = materialize_local_execution_candidates_with_serving(
        planner_state,
        trace_id,
@@ -116,9 +131,13 @@ pub(super) async fn materialize_local_standard_candidate_attempts(
        sticky_session_token.as_deref(),
        input.request_auth_channel.as_deref(),
        persistence_policy,
-        preselection.candidates,
-        preselection.skipped_candidates,
-        LocalCandidateResolutionMode::Standard,
+        candidates,
+        skipped_candidates,
+        if is_gemini_image_bridge {
+            LocalCandidateResolutionMode::WithoutTransportPairGate
+        } else {
+            LocalCandidateResolutionMode::Standard
+        },
        |eligible| {
            let provider_api_format = eligible.provider_api_format.clone();
            let (execution_strategy, conversion_mode) = ai_local_execution_contract_for_formats(
@@ -188,7 +207,18 @@ pub(super) async fn build_local_standard_candidate_attempt_source<'a>(
        input.required_capabilities.as_ref(),
        LocalCandidatePersistencePolicyKind::StandardDecision,
    );
-    Ok(
+    if spec_metadata.api_format == "gemini:generate_content"
+        && gemini_request_is_image_generation(body_json)
+    {
+        let (attempts, candidate_count) =
+            materialize_local_standard_candidate_attempts(state, trace_id, input, body_json, spec)
+                .await?;
+        let source =
+            LocalExecutionCandidateAttemptSource::from_static_attempts_for_image_bridge(attempts);
+        return Ok((source, candidate_count));
+    }
+
+    let (source, candidate_count) =
        build_lazy_requested_model_execution_candidate_attempt_source_with_serving(
            planner_state,
            trace_id,
@@ -253,6 +283,51 @@ pub(super) async fn build_local_standard_candidate_attempt_source<'a>(
                skipped_candidate
            },
        )
-        .await,
+        .await;
+    Ok((source, candidate_count))
+}
+
+#[allow(clippy::too_many_arguments)]
+async fn maybe_append_gemini_image_openai_image_preselection(
+    state: &AppState,
+    planner_state: PlannerAppState<'_>,
+    trace_id: &str,
+    input: &LocalStandardDecisionInput,
+    body_json: &serde_json::Value,
+    spec: LocalStandardSpec,
+    mut candidates: Vec<aether_scheduler_core::SchedulerMinimalCandidateSelectionCandidate>,
+    mut skipped_candidates: Vec<
+        crate::ai_serving::planner::candidate_resolution::SkippedLocalExecutionCandidate,
+    >,
+) -> Result<
+    (
+        Vec<aether_scheduler_core::SchedulerMinimalCandidateSelectionCandidate>,
+        Vec<crate::ai_serving::planner::candidate_resolution::SkippedLocalExecutionCandidate>,
+    ),
+    GatewayError,
+> {
+    let spec_metadata = local_standard_spec_metadata(spec);
+    if spec_metadata.api_format != "gemini:generate_content"
+        || !gemini_request_is_image_generation(body_json)
+    {
+        return Ok((candidates, skipped_candidates));
+    }
+
+    let image_preselection = preselect_local_execution_candidates_for_api_formats_with_serving(
+        planner_state,
+        spec_metadata.api_format,
+        &input.requested_model,
+        spec_metadata.require_streaming,
+        input.required_capabilities.as_ref(),
+        &input.auth_snapshot,
+        input.client_session_affinity.as_ref(),
+        false,
+        LocalCandidatePreselectionKeyMode::ProviderEndpointKeyModelAndApiFormat,
+        vec!["openai:image".to_string()],
    )
+    .await?;
+    candidates.extend(image_preselection.candidates);
+    skipped_candidates.extend(image_preselection.skipped_candidates);
+    let _ = (state, trace_id);
+    Ok((candidates, skipped_candidates))
 }
@@ -18,10 +18,13 @@ use crate::ai_serving::transport::kiro::{
    KIRO_ENVELOPE_NAME,
 };
 use crate::ai_serving::transport::{
-    build_kiro_cross_format_upstream_url, build_standard_provider_request_headers,
-    StandardProviderRequestHeadersInput,
+    build_kiro_cross_format_upstream_url, build_openai_image_headers,
+    build_openai_image_upstream_url, build_standard_provider_request_headers,
+    openai_image_transport_unsupported_reason, resolve_openai_image_auth,
+    ProviderOpenAiImageHeadersInput, StandardProviderRequestHeadersInput,
 };
 use crate::ai_serving::{
+    build_openai_image_request_body_from_gemini_image_request, gemini_request_is_image_generation,
    CandidateFailureDiagnostic, GatewayProviderTransportSnapshot, LocalResolvedOAuthRequestAuth,
 };
 use crate::AppState;
@@ -59,6 +62,15 @@ pub(crate) async fn resolve_local_standard_candidate_payload_parts(
    let candidate = &attempt.eligible.candidate;
    let transport = &attempt.eligible.transport;
    let provider_api_format = attempt.eligible.provider_api_format.as_str();
+    if spec_metadata.api_format == "gemini:generate_content"
+        && provider_api_format == "openai:image"
+        && gemini_request_is_image_generation(body_json)
+    {
+        return resolve_local_gemini_image_to_openai_image_candidate_payload_parts(
+            state, parts, trace_id, body_json, input, attempt,
+        )
+        .await;
+    }
    let is_kiro_claude_cli = is_kiro_claude_messages_transport(transport, provider_api_format);
    let Some(conversion_kind) =
        crate::ai_serving::request_conversion_kind(spec_metadata.api_format, provider_api_format)
@@ -331,6 +343,141 @@ pub(crate) async fn resolve_local_standard_candidate_payload_parts(
    })
 }

+async fn resolve_local_gemini_image_to_openai_image_candidate_payload_parts(
+    state: &AppState,
+    parts: &http::request::Parts,
+    trace_id: &str,
+    body_json: &serde_json::Value,
+    input: &LocalStandardDecisionInput,
+    attempt: &LocalStandardCandidateAttempt,
+) -> Option<LocalStandardCandidatePayloadParts> {
+    let client_api_format = "gemini:generate_content";
+    let provider_api_format = "openai:image";
+    let planner_state = crate::ai_serving::PlannerAppState::new(state);
+    let candidate = &attempt.eligible.candidate;
+    let transport = &attempt.eligible.transport;
+
+    if let Some(skip_reason) =
+        openai_image_transport_unsupported_reason(transport, provider_api_format)
+    {
+        mark_skipped_local_standard_candidate(
+            state,
+            input,
+            trace_id,
+            candidate,
+            attempt.candidate_index,
+            &attempt.candidate_id,
+            skip_reason,
+        )
+        .await;
+        return None;
+    }
+
+    let prepared_candidate = match prepare_header_authenticated_candidate(
+        planner_state,
+        transport,
+        candidate,
+        resolve_openai_image_auth(transport),
+        OauthPreparationContext {
+            trace_id,
+            api_format: provider_api_format,
+            operation: "gemini_image_to_openai_image_candidate_request",
+        },
+    )
+    .await
+    {
+        Ok(prepared) => prepared,
+        Err(skip_reason) => {
+            mark_skipped_local_standard_candidate(
+                state,
+                input,
+                trace_id,
+                candidate,
+                attempt.candidate_index,
+                &attempt.candidate_id,
+                skip_reason,
+            )
+            .await;
+            return None;
+        }
+    };
+
+    let Some(converted) = build_openai_image_request_body_from_gemini_image_request(
+        body_json,
+        parts.uri.path(),
+        &prepared_candidate.mapped_model,
+    ) else {
+        mark_skipped_local_standard_candidate_with_extra_data(
+            state,
+            input,
+            trace_id,
+            candidate,
+            attempt.candidate_index,
+            &attempt.candidate_id,
+            "provider_request_body_build_failed",
+            request_body_build_failure_extra_data(
+                body_json,
+                client_api_format,
+                provider_api_format,
+            ),
+        )
+        .await;
+        return None;
+    };
+
+    let upstream_is_stream = true;
+    let upstream_url = build_openai_image_upstream_url(transport, None);
+    let Some(mut provider_request_headers) =
+        build_openai_image_headers(ProviderOpenAiImageHeadersInput {
+            headers: &parts.headers,
+            auth_header: &prepared_candidate.auth_header,
+            auth_value: &prepared_candidate.auth_value,
+            header_rules: transport.endpoint.header_rules.as_ref(),
+            provider_request_body: &converted.body_json,
+            original_request_body: body_json,
+        })
+    else {
+        mark_skipped_local_standard_candidate_with_failure_diagnostic(
+            state,
+            input,
+            trace_id,
+            candidate,
+            attempt.candidate_index,
+            &attempt.candidate_id,
+            "transport_header_rules_apply_failed",
+            CandidateFailureDiagnostic::header_rules_apply_failed(
+                client_api_format,
+                provider_api_format,
+                "gemini_image_to_openai_image_headers",
+            ),
+        )
+        .await;
+        return None;
+    };
+    apply_codex_openai_responses_special_headers(
+        &mut provider_request_headers,
+        &converted.body_json,
+        &parts.headers,
+        transport.provider.provider_type.as_str(),
+        provider_api_format,
+        Some(trace_id),
+        transport.key.decrypted_auth_config.as_deref(),
+    );
+
+    Some(LocalStandardCandidatePayloadParts {
+        auth_header: prepared_candidate.auth_header,
+        auth_value: prepared_candidate.auth_value,
+        mapped_model: converted.mapped_model,
+        provider_api_format: provider_api_format.to_string(),
+        provider_request_body: converted.body_json,
+        provider_request_headers,
+        upstream_url,
+        upstream_is_stream,
+        envelope_name: None,
+        transport: Arc::clone(transport),
+    })
+}
+
 #[allow(clippy::too_many_arguments)]
 async fn build_kiro_cross_format_payload_parts(
    state: &AppState,