Aether/apps/aether-gateway/src/handlers/admin/observability/monitoring/trace.rs

use super::route_filters::parse_admin_monitoring_limit;
use crate::handlers::admin::request::{AdminAppState, AdminRequestContext};
use crate::log_ids::short_request_id;
use crate::GatewayError;
use aether_admin::observability::monitoring::{
    admin_monitoring_bad_request_response, admin_monitoring_trace_not_found_response,
    admin_monitoring_trace_provider_id_from_path, admin_monitoring_trace_request_id_from_path,
    build_admin_monitoring_trace_provider_stats_payload_response,
    build_admin_monitoring_trace_request_payload_response_with_key_accounts,
    parse_admin_monitoring_attempted_only, AdminMonitoringKeyAccountDisplay,
};
use aether_data_contracts::repository::{
    candidates::{DecisionTrace, RequestCandidateStatus},
    provider_catalog::StoredProviderCatalogKey,
    usage::StoredRequestUsageAudit,
};
use axum::{
    body::Body,
    response::{IntoResponse, Response},
};
use serde_json::{Map, Value};
use std::collections::BTreeMap;
use tracing::debug;

struct ResolvedAdminMonitoringTrace {
    trace: DecisionTrace,
    usage: Option<StoredRequestUsageAudit>,
}

pub(super) async fn build_admin_monitoring_trace_request_response(
    state: &AdminAppState<'_>,
    request_context: &AdminRequestContext<'_>,
) -> Result<Response<Body>, GatewayError> {
    let admin_state = state;
    let Some(request_id) =
        admin_monitoring_trace_request_id_from_path(&request_context.request_path)
    else {
        return Ok(admin_monitoring_bad_request_response("缺少 request_id"));
    };
    let attempted_only = match parse_admin_monitoring_attempted_only(
        request_context.request_query_string.as_deref(),
    ) {
        Ok(value) => value,
        Err(detail) => return Ok(admin_monitoring_bad_request_response(detail)),
    };

    let Some(resolved) =
        resolve_admin_monitoring_trace(admin_state, &request_id, attempted_only).await?
    else {
        debug!(
            event_name = "admin_monitoring_request_trace_not_found",
            log_type = "admin_monitoring",
            request_id = %short_request_id(request_id.as_str()),
            attempted_only,
            path = %request_context.request_path,
            "admin monitoring request trace not found"
        );
        return Ok(admin_monitoring_trace_not_found_response(
            &request_id,
            attempted_only,
        ));
    };
    let key_accounts =
        build_admin_monitoring_key_account_display_map(admin_state, &resolved.trace).await?;

    Ok(
        build_admin_monitoring_trace_request_payload_response_with_key_accounts(
            &resolved.trace,
            resolved.usage.as_ref(),
            &key_accounts,
        ),
    )
}

async fn resolve_admin_monitoring_trace(
    state: &AdminAppState<'_>,
    request_id: &str,
    attempted_only: bool,
) -> Result<Option<ResolvedAdminMonitoringTrace>, GatewayError> {
    let app = state.as_ref();
    if let Some(trace) = app
        .data
        .read_decision_trace(request_id, attempted_only)
        .await
        .map_err(|err| GatewayError::Internal(err.to_string()))?
    {
        let usage = app
            .data
            .read_request_usage_audit(request_id)
            .await
            .map_err(|err| GatewayError::Internal(err.to_string()))?;
        return Ok(Some(ResolvedAdminMonitoringTrace { trace, usage }));
    }

    let mut usage_candidates = Vec::new();
    if let Some(usage) = app
        .data
        .read_request_usage_audit(request_id)
        .await
        .map_err(|err| GatewayError::Internal(err.to_string()))?
    {
        usage_candidates.push(usage);
    }
    if let Some(usage) = state.find_request_usage_by_id(request_id).await? {
        if !usage_candidates.iter().any(|item| item.id == usage.id) {
            usage_candidates.push(usage);
        }
    }

    for usage in usage_candidates {
        for trace_request_id in admin_monitoring_usage_trace_request_ids(&usage) {
            if trace_request_id == request_id {
                continue;
            }
            if let Some(trace) = app
                .data
                .read_decision_trace(&trace_request_id, attempted_only)
                .await
                .map_err(|err| GatewayError::Internal(err.to_string()))?
            {
                return Ok(Some(ResolvedAdminMonitoringTrace {
                    trace,
                    usage: Some(usage),
                }));
            }
        }
    }

    Ok(None)
}

fn admin_monitoring_usage_trace_request_ids(usage: &StoredRequestUsageAudit) -> Vec<String> {
    let mut ids = Vec::new();
    push_non_empty_unique(&mut ids, usage.request_id.as_str());
    if let Some(trace_id) = usage.trace_id() {
        push_non_empty_unique(&mut ids, trace_id);
    }
    if let Some(trace_id) = usage_trace_id_from_headers(usage.request_headers.as_ref()) {
        push_non_empty_unique(&mut ids, trace_id.as_str());
    }
    if let Some(trace_id) = usage_trace_id_from_headers(usage.provider_request_headers.as_ref()) {
        push_non_empty_unique(&mut ids, trace_id.as_str());
    }
    ids
}

fn usage_trace_id_from_headers(headers: Option<&Value>) -> Option<String> {
    let object = headers?.as_object()?;
    object.iter().find_map(|(key, value)| {
        key.eq_ignore_ascii_case(crate::constants::TRACE_ID_HEADER)
            .then(|| {
                value
                    .as_str()
                    .map(str::trim)
                    .filter(|value| !value.is_empty())
            })
            .flatten()
            .map(ToOwned::to_owned)
    })
}

fn push_non_empty_unique(values: &mut Vec<String>, value: &str) {
    let value = value.trim();
    if value.is_empty() || values.iter().any(|existing| existing == value) {
        return;
    }
    values.push(value.to_string());
}

async fn build_admin_monitoring_key_account_display_map(
    state: &AdminAppState<'_>,
    trace: &DecisionTrace,
) -> Result<BTreeMap<String, AdminMonitoringKeyAccountDisplay>, GatewayError> {
    let key_ids = trace
        .candidates
        .iter()
        .filter_map(|item| item.candidate.key_id.as_deref())
        .filter(|value| !value.trim().is_empty())
        .map(ToOwned::to_owned)
        .collect::<std::collections::BTreeSet<_>>()
        .into_iter()
        .collect::<Vec<_>>();
    if key_ids.is_empty() {
        return Ok(BTreeMap::new());
    }

    let keys = state.read_provider_catalog_keys_by_ids(&key_ids).await?;
    Ok(keys
        .into_iter()
        .filter_map(|key| {
            let display = resolve_admin_monitoring_key_account_display(state, &key)?;
            Some((key.id, display))
        })
        .collect())
}

fn resolve_admin_monitoring_key_account_display(
    state: &AdminAppState<'_>,
    key: &StoredProviderCatalogKey,
) -> Option<AdminMonitoringKeyAccountDisplay> {
    let auth_config = parse_admin_monitoring_key_auth_config(state, key);
    let label = auth_config
        .as_ref()
        .and_then(|config| {
            first_non_empty_json_string([
                config.get("email"),
                config.get("account_name"),
                config.get("accountName"),
                config.get("client_email"),
                config.get("account_id"),
                config.get("accountId"),
            ])
        })
        .or_else(|| {
            key.upstream_metadata.as_ref().and_then(|metadata| {
                first_non_empty_json_string([
                    metadata.get("email"),
                    metadata.get("account_name"),
                    metadata.get("accountName"),
                    metadata.get("account_id"),
                    metadata.get("accountId"),
                ])
            })
        });
    let oauth_plan_type = auth_config.as_ref().and_then(|config| {
        first_non_empty_json_string([config.get("plan_type"), config.get("planType")])
    });

    if label.is_none() && oauth_plan_type.is_none() {
        return None;
    }

    Some(AdminMonitoringKeyAccountDisplay {
        label,
        oauth_plan_type,
    })
}

fn parse_admin_monitoring_key_auth_config(
    state: &AdminAppState<'_>,
    key: &StoredProviderCatalogKey,
) -> Option<Map<String, Value>> {
    let ciphertext = key.encrypted_auth_config.as_deref()?;
    let plaintext = state.decrypt_catalog_secret_with_fallbacks(ciphertext)?;
    serde_json::from_str::<Value>(&plaintext)
        .ok()?
        .as_object()
        .cloned()
}

fn first_non_empty_json_string<'a>(
    values: impl IntoIterator<Item = Option<&'a Value>>,
) -> Option<String> {
    values.into_iter().find_map(|value| {
        let text = value?.as_str()?.trim();
        (!text.is_empty()).then(|| text.to_string())
    })
}

pub(super) async fn build_admin_monitoring_trace_provider_stats_response(
    state: &AdminAppState<'_>,
    request_context: &AdminRequestContext<'_>,
) -> Result<Response<Body>, GatewayError> {
    let state = state.as_ref();
    let Some(provider_id) =
        admin_monitoring_trace_provider_id_from_path(&request_context.request_path)
    else {
        return Ok(admin_monitoring_bad_request_response("缺少 provider_id"));
    };
    let limit = match parse_admin_monitoring_limit(request_context.request_query_string.as_deref())
    {
        Ok(value) => value,
        Err(detail) => return Ok(admin_monitoring_bad_request_response(detail)),
    };

    let candidates = state
        .read_request_candidates_by_provider_id(&provider_id, limit)
        .await?;
    let total_attempts = candidates.len();
    let success_count = candidates
        .iter()
        .filter(|item| item.status == RequestCandidateStatus::Success)
        .count();
    let failed_count = candidates
        .iter()
        .filter(|item| item.status == RequestCandidateStatus::Failed)
        .count();
    let cancelled_count = candidates
        .iter()
        .filter(|item| item.status == RequestCandidateStatus::Cancelled)
        .count();
    let skipped_count = candidates
        .iter()
        .filter(|item| item.status == RequestCandidateStatus::Skipped)
        .count();
    let pending_count = candidates
        .iter()
        .filter(|item| item.status == RequestCandidateStatus::Pending)
        .count();
    let available_count = candidates
        .iter()
        .filter(|item| item.status == RequestCandidateStatus::Available)
        .count();
    let unused_count = candidates
        .iter()
        .filter(|item| item.status == RequestCandidateStatus::Unused)
        .count();
    let completed_count = success_count + failed_count;
    let failure_rate = if completed_count == 0 {
        0.0
    } else {
        ((failed_count as f64 / completed_count as f64) * 10000.0).round() / 100.0
    };
    let latency_values = candidates
        .iter()
        .filter_map(|item| item.latency_ms.map(|value| value as f64))
        .collect::<Vec<_>>();
    let avg_latency_ms = if latency_values.is_empty() {
        0.0
    } else {
        let total = latency_values.iter().sum::<f64>();
        ((total / latency_values.len() as f64) * 100.0).round() / 100.0
    };

    Ok(
        build_admin_monitoring_trace_provider_stats_payload_response(
            provider_id,
            total_attempts,
            success_count,
            failed_count,
            cancelled_count,
            skipped_count,
            pending_count,
            available_count,
            unused_count,
            failure_rate,
            avg_latency_ms,
        ),
    )
}